Update index.rst

* Fix

.gitignore

@@ -134,4 +134,10 @@ output/
 autogen/
 
 # cutlass kernel
-!csrc/gpu/cutlass_kernels/gemm/collective/builders
+!csrc/gpu/cutlass_kernels/gemm/collective/builders
+
+
+#fp8
+ops/csrc/fp8/deep_gemm/include/cutlass
+ops/csrc/fp8/deep_gemm/include/cute
+

README.md

@@ -33,16 +33,23 @@
 
 ## News 📢
 
-* **2025.02.20 🔥🔥《PP-UIE信息抽取智能引擎全新升级》** 强化零样本学习能力，支持极少甚至零标注数据实现高效冷启动与迁移学习，显著降低数据标注成本；具备处理长文本能力，支持 8192 个Token长度文档信息抽取，实现跨段落识别关键信息，形成完整理解；提供完整可定制化的训练和推理全流程，训练效率相较于LLama-Factory实现了1.8倍的提升。
-2月26日（周三）19：00为您深度解析全新PP-UIE技术方案及在部署方面的功能、优势与技巧。报名链接：https://www.wjx.top/vm/mBKC6pb.aspx?udsid=606418
+* **2025.03.17 《DeepSeek-R1满血版单机部署实测》** 🔥🔥🔥 飞桨框架3.0大模型推理部署全面升级，支持多款主流大模型，DeepSeek-R1满血版实现单机部署，吞吐提升一倍！欢迎广大用户开箱体验～现已开启有奖活动：完成DeepSeek-R1-MTP 单机部署任务、提交高质量测评blog，即可实时赢取奖金！💰💰💰
+报名[地址](https://www.wjx.top/vm/OlzzmbG.aspx#)， 活动详情：https://github.com/PaddlePaddle/PaddleNLP/issues/10166 ， 参考文档：https://github.com/PaddlePaddle/PaddleNLP/issues/10157 。
+
+* **2025.03.12 [PaddleNLP v3.0 Beta4](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta4)**：全面支持 DeepSeek V3/R1/R1-Distill, 及QwQ-32B等热门思考模型。**DeepSeek V3/R1完整版支持FP8、INT8、4-bit量化推理，MTP投机解码**。单机FP8推理输出超**1000 tokens/s**; 4-bit推理输出超**2100 tokens/s**! 发布新版推理部署镜像，热门模型[一键部署](https://paddlenlp.readthedocs.io/zh/latest/llm/server/docs/general_model_inference.html)。推理部署[使用文档](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/predict/index.html)全面更新，体验全面提升！自研下一代通用信息抽取模型 PP-UIE [全新发布](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/application/information_extraction)，支持8K长度信息抽取。新增大模型 Embedding 训练，支持INF-CL超大batch size训练。新增[MergeKit](https://paddlenlp.readthedocs.io/zh/latest/llm/docs/mergekit.html)模型融合工具，缓解对齐代价。低资源训练全面优化，16G小显存可以流畅训练。
+
+* **2025.03.06 PaddleNLP 现已支持 Qwen/QwQ-32B 模型**: 其模型参数仅有 32B，但其数学推理、编程能力和通用能力可与具备 671B 参数（其中 37B 被激活）的 DeepSeek-R1 媲美。借助 PaddleNLP 3.0套件，现可实现多种并行策略[微调训练](./llm/README.md)、[高性能推理、低比特量化](./llm/docs/predict/qwen.md)和[服务化部署](./llm/server/README.md)。
 
 * **2025.02.10 PaddleNLP 现已支持 DeepSeek-R1系列模型，[在线使用](https://aistudio.baidu.com/projectdetail/8775758)**：依托全新的 PaddleNLP 3.0套件，DeepSeek-R1系列模型现已全面支持。凭借数据并行、数据分组切分并行、模型并行、流水线并行以及专家并行等一系列先进的分布式训练能力，结合 Paddle 框架独有的列稀疏注意力掩码表示技术——FlashMask 方法，DeepSeek-R1系列模型在训练过程中显著降低了显存消耗，同时取得了卓越的训练性能提升。
 
+<details><summary> <b>点击展开</b> </summary><div>
+
+* **2025.02.20 🔥🔥《PP-UIE 信息抽取智能引擎全新升级》** 强化零样本学习能力，支持极少甚至零标注数据实现高效冷启动与迁移学习，显著降低数据标注成本；具备处理长文本能力，支持 8192 个 Token 长度文档信息抽取，实现跨段落识别关键信息，形成完整理解；提供完整可定制化的训练和推理全流程，训练效率相较于 LLama-Factory 实现了1.8倍的提升。
+2月26日（周三）19：00为您深度解析全新 PP-UIE 技术方案及在部署方面的功能、优势与技巧。报名链接：https://www.wjx.top/vm/mBKC6pb.aspx?udsid=606418
+
 * **2024.12.16 [PaddleNLP v3.0 Beta3](https://github.com/PaddlePaddle/PaddleNLP/releases/tag/v3.0.0-beta3)**：大模型功能全新升级，新增了 Llama-3.2、DeepSeekV2模型，升级了 TokenizerFast，快速分词，重构了 SFTTrainer，一键开启 SFT 训练。此外，PaddleNLP 还支持了优化器状态的卸载和重载功能，实现了精细化的重新计算，训练性能提升7%。在 Unified Checkpoint 方面，进一步优化了异步保存逻辑，新增 Checkpoint 压缩功能，可节省78.5%存储空间。
 最后，在大模型推理方面，升级 Append Attention，支持了 FP8量化，支持投机解码。
 
-<details><summary> <b>点击展开</b> </summary><div>
-
 * **2024.12.13 📚《飞桨大模型套件 Unified Checkpoint 技术》**，加速模型存储95%，节省空间78%。支持全分布式策略调整自适应转换，提升模型训练的灵活性与可扩展性。训练-压缩-推理统一存储协议，无需手动转换提升全流程体验。Checkpoint 无损压缩结合异步保存，实现秒级存储并降低模型存储成本。适用于智能制造、指挥交通、医疗健康、金融服务等产业实际场景。12月24日（周二）19：00直播为您详细解读该技术如何优化大模型训练流程。报名链接：https://www.wjx.top/vm/huZkHn9.aspx?udsid=787976
 
 * **2024.11.28 📚《FlashRAG-Paddle | 基于 PaddleNLP 的高效开发与评测 RAG 框架》**，为文本更快更好构建准确嵌入表示、加速推理生成速度。PaddleNLP 支持超大 Batch 嵌入表示学习与多硬件高性能推理，涵盖 INT8/INT4量化技术及多种高效注意力机制优化与 TensorCore 深度优化。内置全环节算子融合技术，使得 FlashRAG 推理性能相比 transformers 动态图提升70%以上，结合检索增强知识输出结果更加准确，带来敏捷高效的使用体验。直播时间：12月3日（周二）19：00。报名链接：https://www.wjx.top/vm/eaBa1vA.aspx?udsid=682361
@@ -109,14 +116,14 @@
 |           [Qwen1.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)            | Qwen/Qwen1.5-0.5B, Qwen/Qwen1.5-0.5B-Chat, Qwen/Qwen1.5-1.8B, Qwen/Qwen1.5-1.8B-Chat, Qwen/Qwen1.5-4B, Qwen/Qwen1.5-4B-Chat, Qwen/Qwen1.5-7B, Qwen/Qwen1.5-7B-Chat, Qwen/Qwen1.5-14B, Qwen/Qwen1.5-14B-Chat, Qwen/Qwen1.5-32B, Qwen/Qwen1.5-32B-Chat, Qwen/Qwen1.5-72B, Qwen/Qwen1.5-72B-Chat, Qwen/Qwen1.5-110B, Qwen/Qwen1.5-110B-Chat, Qwen/Qwen1.5-MoE-A2.7B, Qwen/Qwen1.5-MoE-A2.7B-Chat |
 |            [Qwen2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)             | Qwen/Qwen2-0.5B, Qwen/Qwen2-0.5B-Instruct, Qwen/Qwen2-1.5B, Qwen/Qwen2-1.5B-Instruct, Qwen/Qwen2-7B, Qwen/Qwen2-7B-Instruct, Qwen/Qwen2-72B, Qwen/Qwen2-72B-Instruct, Qwen/Qwen2-57B-A14B, Qwen/Qwen2-57B-A14B-Instruct                                                                                                                                                                       |
 |          [Qwen2-Math](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)          | Qwen/Qwen2-Math-1.5B, Qwen/Qwen2-Math-1.5B-Instruct, Qwen/Qwen2-Math-7B, Qwen/Qwen2-Math-7B-Instruct, Qwen/Qwen2-Math-72B, Qwen/Qwen2-Math-72B-Instruct, Qwen/Qwen2-Math-RM-72B                                                                                                                                                                                                               |
-|           [Qwen2.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)            | Qwen/Qwen2.5-0.5B, Qwen/Qwen2.5-0.5B-Instruct, Qwen/Qwen2.5-1.5B, Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-3B, Qwen/Qwen2.5-3B-Instruct, Qwen/Qwen2.5-7B, Qwen/Qwen2.5-7B-Instruct, Qwen/Qwen2.5-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, Qwen/Qwen2.5-72B, Qwen/Qwen2.5-72B-Instruct                                                                     |
+|           [Qwen2.5](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)            | Qwen/Qwen2.5-0.5B, Qwen/Qwen2.5-0.5B-Instruct, Qwen/Qwen2.5-1.5B, Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-3B, Qwen/Qwen2.5-3B-Instruct, Qwen/Qwen2.5-7B, Qwen/Qwen2.5-7B-Instruct, Qwen/Qwen2.5-7B-Instruct-1M, Qwen/Qwen2.5-14B, Qwen/Qwen2.5-14B-Instruct, Qwen/Qwen2.5-14B-Instruct-1M, Qwen/Qwen2.5-32B, Qwen/Qwen2.5-32B-Instruct, Qwen/Qwen2.5-72B, Qwen/Qwen2.5-72B-Instruct          |
 |         [Qwen2.5-Math](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)         | Qwen/Qwen2.5-Math-1.5B, Qwen/Qwen2.5-Math-1.5B-Instruct, Qwen/Qwen2.5-Math-7B, Qwen/Qwen2.5-Math-7B-Instruct, Qwen/Qwen2.5-Math-72B, Qwen/Qwen2.5-Math-72B-Instruct, Qwen/Qwen2.5-Math-RM-72B                                                                                                                                                                                                 |
 |        [Qwen2.5-Coder](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)         | Qwen/Qwen2.5-Coder-1.5B, Qwen/Qwen2.5-Coder-1.5B-Instruct, Qwen/Qwen2.5-Coder-7B, Qwen/Qwen2.5-Coder-7B-Instruct                                                                                                                                                                                                                                                                              |
+|             [QwQ](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/qwen/)              | Qwen/QwQ-32B, Qwen/QwQ-32B-Preview                                                                                                                                                                                                                                                                                                                                                            |
 |            [Yuan2](https://github.com/PaddlePaddle/PaddleNLP/tree/develop/llm/config/yuan/)             | IEITYuan/Yuan2-2B, IEITYuan/Yuan2-51B, IEITYuan/Yuan2-102B                                                                                                                                                                                                                                                                                                                                    |
 
 * 4D 并行和算子优化已支持 LLaMA 系列、Baichuan 系列、Bloom 系列、ChatGLM 系列、Gemma 系列、Mistral 系列、OPT 系列和 Qwen 系列，【LLM】模型4D 并行和算子支持列表如下：
 
-
 | 模型名称/并行能力支持 | 数据并行 | 张量模型并行 |          | 参数分片并行 |        |        | 流水线并行 |
 |:---------------------:|:--------:|:------------:|:--------:|:------------:|:------:|:------:|:----------:|
 |                       |          |   基础能力   | 序列并行 |    stage1    | stage2 | stage3 |            |
@@ -154,31 +161,32 @@
 | [OPT](./llm/config/opt)                    |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         🚧         |  🚧  |    ✅     |      🚧      |
 | [Gemma](./llm/config/gemma)                |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         ✅          |  🚧  |    ✅     |      🚧      |
 | [Yuan](./llm/config/yuan)                  |    ✅     |  ✅  |  ✅   |    🚧     |      🚧       |         ✅          |  🚧  |    ✅     |      🚧      |
-* [大模型推理](./llm/docs/predict/inference.md)已支持 LLaMA 系列、Qwen 系列、Mistral 系列、ChatGLM 系列、Bloom 系列和 Baichuan 系列，支持 Weight Only INT8及 INT4推理，支持 WAC（权重、激活、Cache KV）进行 INT8、FP8量化的推理，【LLM】模型推理支持列表如下：
+* [大模型推理](./llm/docs/predict/inference.md)已支持 LLaMA 系列、Qwen 系列、DeepSeek 系列、Mistral 系列、ChatGLM 系列、Bloom 系列和 Baichuan 系列，支持 Weight Only INT8及 INT4推理，支持 WAC（权重、激活、Cache KV）进行 INT8、FP8量化的推理，【LLM】模型推理支持列表如下：
 
-|          模型名称/量化类型支持           | FP16/BF16 | WINT8 | WINT4 | INT8-A8W8 | FP8-A8W8 | INT8-A8W8C8 |
-|:----------------------------------------:|:---------:|:-----:|:-----:|:---------:|:--------:|:-----------:|
-|   [LLaMA](./llm/docs/predict/llama.md)   |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
-|    [Qwen](./llm/docs/predict/qwen.md)    |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
-|  [Qwen-Moe](./llm/docs/predict/qwen.md)  |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
-| [Mixtral](./llm/docs/predict/mixtral.md) |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
-|                 ChatGLM                  |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
-|                  Bloom                   |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
-|                 BaiChuan                 |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |     🚧      |
+|           模型名称/量化类型支持            | FP16/BF16 | WINT8 | WINT4 | INT8-A8W8 | FP8-A8W8 | INT8-A8W8C8 |
+|:------------------------------------------:|:---------:|:-----:|:-----:|:---------:|:--------:|:-----------:|
+|    [LLaMA](./llm/docs/predict/llama.md)    |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
+|     [Qwen](./llm/docs/predict/qwen.md)     |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |      ✅      |
+| [DeepSeek](./llm/docs/predict/deepseek.md) |     ✅     |   ✅   |   ✅   |    🚧     |    ✅     |     🚧      |
+|   [Qwen-Moe](./llm/docs/predict/qwen.md)   |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|  [Mixtral](./llm/docs/predict/mixtral.md)  |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                  ChatGLM                   |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                   Bloom                    |     ✅     |   ✅   |   ✅   |    🚧     |    🚧    |     🚧      |
+|                  BaiChuan                  |     ✅     |   ✅   |   ✅   |     ✅     |    ✅     |     🚧      |
 
 ## 安装
 
 ### 环境依赖
 
 * python >= 3.8
-* paddlepaddle >= 3.0.0rc0
+* paddlepaddle >= 3.0.0rc1
 
 如果您尚未安装 PaddlePaddle，请参考 [飞桨官网](https://www.paddlepaddle.org.cn/) 进行安装。
 
 ### pip 安装
 
 ```shell
-pip install --upgrade paddlenlp==3.0.0b3
+pip install --upgrade paddlenlp==3.0.0b4
 ```
 
 或者可通过以下命令安装最新 develop 分支代码：

csrc/README.md

@@ -23,5 +23,6 @@ python setup_cuda.py install
 - 对于89架构的 GPU，CUDA 版本至少为12.4
 - 对于90架构的 GPU，CUDA 版本至少为12.0
 ```shell
+cd tools
 sh tune_fp8_gemm.sh
 ```

csrc/cpu/setup.sh

@@ -47,9 +47,11 @@ fi
 cd xFasterTransformer
 git apply paddle.patch
 
-# #4. build xFasterTransformer
-sh ./3rdparty/prepare_oneccl.sh
-source ./3rdparty/oneccl/build/_install/env/setvars.sh
+#4. build xFasterTransformer
+cd 3rdparty
+bash prepare_oneccl.sh
+source ./oneccl/build/_install/env/setvars.sh
+cd ..
 
 rm -rf build
 mkdir build && cd build

csrc/cpu/src/avx_weight_only.cc

@@ -15,7 +15,7 @@
 #include "matmul_helper.h"
 #include "my_types.h"
 #include "paddle/extension.h"
-git adtemplate <typename T>
+template <typename T>
 void AvxCompute(const paddle::Tensor &x,
                 const paddle::Tensor &weight,
                 bool trans,
@@ -137,12 +137,12 @@ std::vector<paddle::Tensor> InvokeAvxWeightOnly(const paddle::Tensor &x,
 
 std::vector<std::vector<int64_t>> AvxWeightOnlyInferShape(
     std::vector<int64_t> x_shape,
-    std::vector<int64_t> weigh_shape) {
+    std::vector<int64_t> weight_shape) {
   int m = 1;
   for (int i = 0; i < x_shape.size() - 1; i++) {
     m = m * x_shape[i];
   }
-  return {std::vector<int64_t>{m, weigh_shape[1]}};
+  return {std::vector<int64_t>{m, weight_shape[1]}};
 }
 
 std::vector<paddle::DataType> AvxWeightOnlyInferDtype(

csrc/cpu/src/stop_generation_multi_ends.cc

@@ -14,9 +14,20 @@
 
 #include <stdlib.h>
 #include <string.h>
-
-#include "helper.h"
 #include <stdio.h>
+#include <vector>
+#include "paddle/extension.h"
+
+
+bool is_in_end(const int64_t id, const int64_t *end_ids, int length) {
+    bool flag = false;
+    for (int i = 0; i < length; i++) {
+        if (id == end_ids[i]) {
+            return true;
+        }
+    }
+    return flag;
+}
 
 void set_value_by_flags(const bool* stop_flags,
                         const int64_t* end_ids,

csrc/gpu/append_attn/append_attention_kernel.h

@@ -300,10 +300,3 @@ void CascadeAppendAttentionKernel(
         "cache_int4_zp]");
   }
 }
-
-inline uint32_t get_max_partition_size(int bsz) {
-    static const char* max_partition_size_env = std::getenv("FLAGS_cascade_attention_max_partition_size");
-    static const uint32_t max_partition_size =
-            max_partition_size_env == nullptr ? 0 : std::stoul(std::string(max_partition_size_env));
-    return (max_partition_size != 0 ? max_partition_size : (bsz == 1 ? 128 : 512));
-}

csrc/gpu/append_attn/decode_attention_kernel.cu

@@ -28,7 +28,6 @@ do
         exit(1);                                                                \
     }                                                                           \
 }while(0)
-// #define DEBUG_DEC_ATTN
 
 template <typename T, typename OutT, int vec_size, uint32_t bdy, uint32_t HEAD_DIM>
 __global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi_out, // [bsz, num_chunks, num_heads, head_dim]
@@ -94,7 +93,6 @@ __global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi
     offset = (qid * num_chunks * num_heads + i * num_heads + hid) * head_dim + vid * vec_size;
     Load<T, vec_size>(&multi_out[offset], &load_vec);
     const T scale1 = hexp(m_prev - m), scale2 = hexp(m_now - m);
-    // const T scale1 = __expf(m_prev - m), scale2 = __expf(m_now - m);
     d = d * scale1 + d_now * scale2;
 #pragma once
     for (int j = 0; j < vec_size; j++) {
@@ -127,102 +125,6 @@ __global__ void merge_varlen_multi_chunks_v2_kernel(const T * __restrict__ multi
   Store<OutT, vec_size>(out_vec, &out[(start_token_ids * num_heads + hid) * head_dim + vid * vec_size]);
 }
 
-// template <typename T, typename OutT, int vec_size>
-// __global__ void merge_varlen_multi_chunks_kernel(const T * __restrict__ multi_out, // [bsz, num_chunks, num_heads, head_dim]
-//                                                 const T * __restrict__ multi_m, // [bsz, num_chunks, num_heads]
-//                                                 const T * __restrict__ multi_d, // [bsz, num_chunks, num_heads]
-//                                                 const int * __restrict__ seq_lens_q,
-//                                                 const int * __restrict__ seq_lens_kv,
-//                                                 const int * __restrict__ cum_offsets,
-//                                                 const T * __restrict__ shift_bias, // [q_num_heads * HEAD_DIM]
-//                                                 const T * __restrict__ smooth_weight, // [q_num_heads * HEAD_DIM]
-//                                                 OutT * __restrict__ out, // [token_num, num_heads, head_dim]
-//                                                 const float in_scale,
-//                                                 const int num_chunks,
-//                                                 const int chunk_size,
-//                                                 const int max_seq_len,
-//                                                 const int num_heads,
-//                                                 const int head_dim) {
-//   const int vid = threadIdx.x, hid = threadIdx.y;
-//   const int qid = blockIdx.x;
-//   const int seq_len_q = seq_lens_q[qid];
-//   if (seq_len_q == 0) return;
-//   int seq_len_kv = seq_lens_kv[qid];
-//   if (seq_len_kv == 0) return;
-//   seq_len_kv += seq_len_q;
-//   const int num_chunks_this_seq = div_up(seq_len_kv, chunk_size);
-//   if (num_chunks_this_seq == 1) {
-//     return;
-//   }
-//   const int start_token_ids = qid * max_seq_len - __ldg(&cum_offsets[qid]);
-//   using LoadT = AlignedVector<T, vec_size>;
-//   LoadT load_vec;
-//   LoadT res_vec;
-//   if constexpr (std::is_same<T, half>::value) {
-// #pragma unroll
-//     for (int i = 0; i < vec_size / 2; ++i) {
-//       *((half2*)(&res_vec) + i) = make_half2(0, 0);
-//     }
-//   } else if constexpr (std::is_same<T, nv_bfloat16>::value) {
-// #pragma unroll
-//     for (int i = 0; i < vec_size / 2; ++i) {
-//       *((nv_bfloat162*)(&res_vec) + i) = make_bfloat162(0, 0);
-//     }
-//   }
-//   T m;
-//   T d = 1.f;
-//   if constexpr (std::is_same<T, half>::value) {
-//     m = __float2half(-5e4f);
-//   } else if constexpr (std::is_same<T, nv_bfloat16>::value) {
-//     m = __float2bfloat16(-3.38953e38f);
-//   }
-// #pragma unroll 2
-//   for (int i = 0; i < num_chunks_this_seq; ++i) {
-//     uint32_t offset = (qid * num_chunks + i) * num_heads + hid;
-//     T m_prev = m;
-//     T d_prev = d;
-//     const T m_now = multi_m[offset];
-//     const T d_now = multi_d[offset];
-//     m = m_prev > m_now ? m_prev : m_now;
-//     offset = (qid * num_chunks * num_heads + i * num_heads + hid) * head_dim + vid * vec_size;
-//     Load<T, vec_size>(&multi_out[offset], &load_vec);
-//     const T scale1 = hexp(m_prev - m), scale2 = hexp(m_now - m);
-//     // const T scale1 = __expf(m_prev - m), scale2 = __expf(m_now - m);
-//     d = d * scale1 + d_now * scale2;
-// #pragma unroll
-//     for (int j = 0; j < vec_size; j++) {
-//       res_vec[j] = res_vec[j] * scale1 + load_vec[j] * scale2;
-//     }
-//   }
-// #pragma unroll 
-//   for (int j = 0; j < vec_size; j++) {
-//     res_vec[j] /= d;
-//   }
-
-//   AlignedVector<OutT, vec_size> out_vec;
-//   if (in_scale > 0) {
-//     const uint32_t shift_smooth_offset = hid * head_dim + vid * vec_size;
-//     AlignedVector<T, vec_size> shift_bias_vec;
-//     AlignedVector<T, vec_size> smooth_weight_vec;
-//     Load<T, vec_size>(shift_bias + shift_smooth_offset, &shift_bias_vec);
-//     Load<T, vec_size>(smooth_weight + shift_smooth_offset, &smooth_weight_vec);
-// #pragma unroll
-//     for (int i = 0; i < vec_size; ++i) {
-//       float quant_value  = 127.0f * static_cast<float>((res_vec[i] + shift_bias_vec[i]) * smooth_weight_vec[i]) * in_scale;
-//       quant_value = rintf(quant_value);
-//       quant_value = quant_value > 127.0f ? 127.0f : quant_value;
-//       quant_value = quant_value < -127.0f ? -127.0f : quant_value;
-//       out_vec[i] = static_cast<OutT>(quant_value);
-//     }
-//   } else {
-// #pragma unroll
-//     for (int i = 0; i < vec_size; ++i) {
-//       out_vec[i] = static_cast<OutT>(res_vec[i]);
-//     }
-//   }
-//   Store<OutT, vec_size>(out_vec, &out[(start_token_ids * num_heads + hid) * head_dim + vid * vec_size]);
-// }
-
 template <bool partition_kv, typename T, typename OutT, typename CacheT, uint32_t NUM_STAGES, uint32_t DEAL_EACH_TIME, uint32_t GROUP_SIZE, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_V, 
           uint32_t BLOCK_SIZE, uint32_t VEC_SIZE, uint32_t CACHE_VEC_SIZE, uint32_t bdx, uint32_t bdy>
 __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [token_num, num_heads, head_dim]
@@ -244,8 +146,11 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
                                                     T * __restrict__ tmp_m, // [batch_size, num_chunks, num_heads]
                                                     T * __restrict__ tmp_d, // [batch_size, num_chunks, num_heads]
                                                     OutT * __restrict__ out) {
+  constexpr uint32_t q_block_num = GROUP_SIZE / bdy;
   const uint32_t bidx = blockIdx.x, kv_head_idx = blockIdx.z;
-  const uint32_t bid = bidx, gid = threadIdx.y;
+  const uint32_t bid = bidx / q_block_num, q_block_idx = bidx % q_block_num;
+  const uint32_t bidy = threadIdx.y;
+  const uint32_t gid = q_block_idx * bdy + threadIdx.y;
   const uint32_t tidx = threadIdx.x;
   constexpr uint32_t num_vec_per_head_qk = HEAD_DIM_QK / VEC_SIZE;
   constexpr uint32_t num_vec_per_head_v = HEAD_DIM_V / VEC_SIZE;
@@ -270,7 +175,6 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   kv_len += q_len;
   const uint32_t num_chunk_this_seq = div_up(kv_len, chunk_size);
   const uint32_t q_start_idx = bid * max_seq_len - __ldg(&cum_offsets[bid]);
-  // const uint32_t q_start_idx = bid;
   const uint32_t q_write_idx = bid * max_seq_len - __ldg(&cum_offsets[bid]);
   if (chunk_id >= num_chunk_this_seq) {
     return;
@@ -283,11 +187,10 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   extern __shared__ uint8_t smem[];
   const T *q_now = q + (q_start_idx * q_num_heads + q_head_idx) * HEAD_DIM_QK;
   T *q_smem = reinterpret_cast<T*>(smem); // [HEAD_DIM_QK * sizeof(T)]
-  T *cu_q_smem = q_smem + gid * HEAD_DIM_QK;
+  T *cu_q_smem = q_smem + bidy * HEAD_DIM_QK;
 #pragma unroll
   for(uint32_t vid = tidx; vid < num_vec_per_head_qk; vid += bdx) {
     ((float4*)(&cu_q_smem[vid * VEC_SIZE]))[0] = ((float4*)(&q_now[vid * VEC_SIZE]))[0];
-
   }
   __syncthreads();
   using VecT = AlignedVector<T, VEC_SIZE>;
@@ -302,14 +205,14 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
   }
 
 
-  CacheT *kv_smem = reinterpret_cast<CacheT*>(smem + GROUP_SIZE * HEAD_DIM_QK * sizeof(CacheT)); // [NUM_STAGES * DEAL_EACH_TIME * HEAD_DIM_QK]
+  CacheT *kv_smem = reinterpret_cast<CacheT*>(smem + bdy * HEAD_DIM_QK * sizeof(CacheT)); 
   uint32_t stage_idx = 0;  
   constexpr int loop_times = DEAL_EACH_TIME / bdy;
 #pragma unroll
   for (int i = 0; i < NUM_STAGES; ++i) {
 #pragma unroll
     for (int j = 0; j < loop_times; ++j) {
-      const uint32_t k_seq_offset = i * DEAL_EACH_TIME + j * bdy + gid;
+      const uint32_t k_seq_offset = i * DEAL_EACH_TIME + j * bdy + bidy;
       const uint32_t k_seq_id = chunk_start + k_seq_offset;
       produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
         kv_smem,
@@ -366,7 +269,7 @@ __global__ void multi_query_decode_attention_kernel(T * __restrict__ q, // [toke
 
 #pragma unroll
     for (int j = 0; j < loop_times; ++j) {
-      const uint32_t k_seq_offset = j * bdy + gid;
+      const uint32_t k_seq_offset = j * bdy + bidy;
       produce_kv<SharedMemFillMode::kNoFill, HEAD_DIM_QK, VEC_SIZE, num_vec_per_head_qk, bdx, BLOCK_SIZE, CACHE_VEC_SIZE>(
         kv_smem,
         cache_k,
@@ -433,10 +336,6 @@ void MultiQueryDecoderAttention(
   auto token_num = meta_data.token_nums;
   auto bsz = meta_data.batch_size;
   auto max_block_num_per_seq = meta_data.max_blocks_per_seq;
-  // const float scale = 1.f / sqrt(HEAD_DIM);
-  // std::cout << "bsz: " << bsz << ", token_num: " << token_num << ", num_heads: " << num_heads << ", kv_num_heads: " << kv_num_heads << ", group_size: " << GROUP_SIZE << std::endl;
-  // dev_ctx.template Alloc<T>(out);
-
   constexpr int num_stages = NUM_STAGE;
 
   constexpr int vec_size = 16 / sizeof(T); // 8 16 32
@@ -444,15 +343,11 @@ void MultiQueryDecoderAttention(
   constexpr int blockxc = HEAD_DIM_QK / cache_vec_size;
   constexpr int num_vec_per_head = HEAD_DIM_QK / vec_size;
   constexpr int blockx = num_vec_per_head < 32 ? num_vec_per_head : 32;
-  // constexpr int blockx = 32;
 
-  constexpr int blocky = GROUP_SIZE;
-  const int gridx = bsz;
-  // static_assert(blockx <= 32);
-  
-  constexpr int num_threads = blockx * blocky;
-  // std::cout << "blockx: " << blockx << ", blocky: " << blocky << ", << ", blockxc: " << blockxc << ", DEAL_EACH_TIME: " << DEAL_EACH_TIME << std::endl;
+  constexpr int blocky = GROUP_SIZE > 16 ? 16 : GROUP_SIZE;
+  const int gridx = bsz * (GROUP_SIZE / blocky);
   
+  constexpr int num_threads = blockx * blocky; 
 
   auto splitkv_kernel = multi_query_decode_attention_kernel<true, NV_TYPE, NV_TYPE, NV_TYPE, num_stages, DEAL_EACH_TIME, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V,
                                                                         BLOCK_SIZE, vec_size, cache_vec_size, blockx, blocky>;
@@ -464,11 +359,8 @@ void MultiQueryDecoderAttention(
   
   const uint32_t chunk_size = get_max_partition_size(bsz);
   const int num_chunks = div_up(max_dec_len, chunk_size);
-  // size_t smem_size = blocky * sizeof(T) * 2 + div_up(max_block_num_per_seq, 4) * 4 * sizeof(int)
-  //                    + cache_smem_bytes + blocky * HEAD_DIM * sizeof(T);
-  size_t smem_size = cache_smem_bytes + GROUP_SIZE * HEAD_DIM_QK * sizeof(T);
-  // size_t smem_size = max(size_t(cache_smem_bytes), HEAD_DIM_QK * sizeof(T)) + blocky * sizeof(T) * 2 + blocky * sizeof(T) * 2 + blocky * sizeof(T) * 2;
-  
+  size_t smem_size = cache_smem_bytes + blocky * HEAD_DIM_QK * sizeof(T);
+
   if (smem_size >= 48 * 1024) {
     cudaFuncSetAttribute(
       splitkv_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
@@ -479,25 +371,16 @@ void MultiQueryDecoderAttention(
   cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
   cudaOccupancyMaxActiveBlocksPerMultiprocessor(
     &act_blocks_per_sm, splitkv_kernel, num_threads, smem_size);
-  // assert(act_blocks_per_sm > 1);
+  assert(act_blocks_per_sm > 1);
   
   const int num_blocks_per_wave = sm_count * act_blocks_per_sm;
   const int num_blocks_need = gridx * num_chunks * kv_num_heads;
   const int max_num_chunks = div_up(num_blocks_per_wave, num_blocks_need);
   const float ratio = static_cast<float>(num_blocks_need) / static_cast<float>(num_blocks_per_wave);
 
-  
-  // std::cout << "num_blocks_per_wave: " << num_blocks_per_wave << ", num_blocks_need: " << num_blocks_need << ", max_num_chunks: " << max_num_chunks << ", ratio: " << ratio << std::endl;
-  // std::cout << "num_chunks: " << num_chunks;
-  
   dim3 grids(gridx, num_chunks, kv_num_heads);
   dim3 blocks(blockx, blocky);
-  // std::cout << "grids: " << grids.x << ", " << grids.y << ", " << grids.z
-  //   << " blocks: " << blocks.x << ", " << blocks.y
-  //   << " smem_size: " << div_up(smem_size, 1024) << "KB" << " HEAD_DIM_QK:" << HEAD_DIM_QK << " HEAD_DIM_V:" << HEAD_DIM_V
-  //   << " num_blocks_need: " << num_blocks_need << ", num_blocks_per_wave: " << num_blocks_per_wave << std::endl;
   if (num_chunks <= 1) {
-    // std::cout << "not split kv";
     auto no_splitkv_kernel = multi_query_decode_attention_kernel<false, NV_TYPE, NV_TYPE, NV_TYPE, num_stages, DEAL_EACH_TIME, GROUP_SIZE, HEAD_DIM_QK, HEAD_DIM_V, BLOCK_SIZE, vec_size, 
                                                                              cache_vec_size, blockx, blocky>;
     if (smem_size >= 48 * 1024) {
@@ -526,10 +409,9 @@ void MultiQueryDecoderAttention(
       reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>()))
     );
 
-    // CHECK(cudaGetLastError());  // 捕捉同步前的最后一个错误。
+    // CHECK(cudaGetLastError());
     // CHECK(cudaDeviceSynchronize());
   } else {
-    // std::cout << "split kv";
     auto *allocator = paddle::GetAllocator(q.place());
     phi::Allocator::AllocationPtr tmp_workspace, tmp_m, tmp_d;
     tmp_workspace = allocator->Allocate(
@@ -541,8 +423,6 @@ void MultiQueryDecoderAttention(
     tmp_d = allocator->Allocate(
         phi::SizeOf(q.dtype()) *
         static_cast<size_t>(bsz * num_chunks * num_heads));
-    // std::cout << "grids: " << grids.x << ", " << grids.y << ", " << grids.z
-    //   << " blocks: " << blocks.x << ", " << blocks.y << std::endl;
 
     splitkv_kernel<<<grids, blocks, smem_size, stream>>>(
       reinterpret_cast<NV_TYPE*>(const_cast<T*>(q.data<T>())),
@@ -565,37 +445,13 @@ void MultiQueryDecoderAttention(
       reinterpret_cast<NV_TYPE*>(tmp_d->ptr()),
       reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>()))
     );
-    // CHECK(cudaGetLastError());  // 捕捉同步前的最后一个错误。
+    // CHECK(cudaGetLastError());
     // CHECK(cudaDeviceSynchronize());
-    // merge
-    // if (num_chunks <= bsz) {
-    //   dim3 grids_merge(bsz);
-    //   dim3 blocks_merge(blockx, num_heads);
-    //   assert(blockx * num_heads <= 1024);
-    //   merge_varlen_multi_chunks_kernel<NV_TYPE, NV_TYPE, vec_size><<<grids_merge, blocks_merge, 0, stream>>>(
-    //     reinterpret_cast<NV_TYPE*>(tmp_workspace->ptr()),
-    //     reinterpret_cast<NV_TYPE*>(tmp_m->ptr()),
-    //     reinterpret_cast<NV_TYPE*>(tmp_d->ptr()),
-    //     seq_lens_q.data<int>(),
-    //     seq_lens_kv.data<int>(),
-    //     cum_offsets.data<int>(),
-    //     reinterpret_cast<NV_TYPE*>(const_cast<T*>(shift_bias_ptr)),
-    //     reinterpret_cast<NV_TYPE*>(const_cast<T*>(smooth_weight_ptr)),
-    //     reinterpret_cast<NV_TYPE*>(const_cast<T*>(out->data<T>())),
-    //     in_scale,
-    //     num_chunks,
-    //     chunk_size,
-    //     max_seq_len,
-    //     num_heads,
-    //     HEAD_DIM_V
-    //   );
-    // } else {
+    
     constexpr int mblockx = HEAD_DIM_V / vec_size;
-    // constexpr int blockx = 32;
     constexpr int bdy = 256 / mblockx;
     dim3 grids_merge(bsz, num_heads);
     dim3 blocks_merge(mblockx, bdy);
-    // std::cout << "work_space size:" << bsz * num_chunks * num_heads * HEAD_DIM_V << std::endl;
     merge_varlen_multi_chunks_v2_kernel<NV_TYPE, NV_TYPE, vec_size, bdy, HEAD_DIM_V><<<grids_merge, blocks_merge, 0, stream>>>(
       reinterpret_cast<NV_TYPE*>(tmp_workspace->ptr()),
       reinterpret_cast<NV_TYPE*>(tmp_m->ptr()),
@@ -614,7 +470,7 @@ void MultiQueryDecoderAttention(
       HEAD_DIM_V
     );
   }
-  // CHECK(cudaGetLastError());  // 捕捉同步前的最后一个错误。
+  // CHECK(cudaGetLastError());
   // CHECK(cudaDeviceSynchronize());
 }
 
@@ -639,7 +495,6 @@ void DecodeMLAAttentionKernel(
   bool causal,
   cudaStream_t &stream,
   paddle::Tensor *out) { 
-  // std::cout << "start: " << std::endl;
   const auto token_num = meta_data.token_nums;
   const auto block_size = meta_data.block_size;
   const auto bsz = meta_data.batch_size;
@@ -653,16 +508,6 @@ void DecodeMLAAttentionKernel(
   const uint32_t num_stage = get_cascade_attention_num_stages();
   const uint32_t num_threads = get_cascade_attention_num_threads();
 
-  // uint32_t cache_type = 0;
-  // if (cache_k_scale) {
-  //   if (cache_k_zp) {
-  //     cache_type = 2;
-  //   } else {
-  //     cache_type = 1;
-  //   }
-  // }
-  // std::cout << "cache_type: " << cache_type << ", group_size: " << group_size << std::endl;
-  // std::cout << "cache_type: " << cache_type << ", group_size: " << group_size << ", deal_each_time: " << deal_each_time << ", num_stage: " << num_stage;
   DISPATCH_CAUSAL(causal, CAUSAL,
     {DISPATCH_MLA_GROUP_SIZE(group_size, GROUP_SIZE,
       {DISPATCH_MLA_HEAD_DIM(head_dim_qk, HEAD_DIM_QK,  
@@ -695,23 +540,23 @@ template void DecodeMLAAttentionKernel<paddle::bfloat16>(
   cudaStream_t &stream,
   paddle::Tensor *out);
 
-// template void DecodeMLAAttentionKernel<paddle::float16>(
-//   const AppendAttnMetaData& meta_data,
-//   const paddle::Tensor &q, // [token_num, num_heads, head_dim]
-//   const paddle::Tensor &cache_k,
-//   const paddle::Tensor &cache_v,
-//   const paddle::optional<paddle::Tensor>& attn_mask,
-//   const paddle::optional<paddle::Tensor>& shift_bias,
-//   const paddle::optional<paddle::Tensor>& smooth_weight,
-//   const paddle::Tensor &seq_lens_q, // q_seq_len is 1
-//   const paddle::Tensor &seq_lens_kv,
-//   const paddle::Tensor &padding_offsets,
-//   const paddle::Tensor &cum_offsets,
-//   const paddle::Tensor &block_table,
-//   int max_seq_len,
-//   int max_dec_len,
-//  float softmax_scale,
-//   float in_scale,
-//   bool causal,
-//   cudaStream_t &stream,
-//   paddle::Tensor *out);
+template void DecodeMLAAttentionKernel<paddle::float16>(
+  const AppendAttnMetaData& meta_data,
+  const paddle::Tensor &q, // [token_num, num_heads, head_dim]
+  const paddle::Tensor &cache_k,
+  const paddle::Tensor &cache_v,
+  const paddle::optional<paddle::Tensor>& attn_mask,
+  const paddle::optional<paddle::Tensor>& shift_bias,
+  const paddle::optional<paddle::Tensor>& smooth_weight,
+  const paddle::Tensor &seq_lens_q, // q_seq_len is 1
+  const paddle::Tensor &seq_lens_kv,
+  const paddle::Tensor &padding_offsets,
+  const paddle::Tensor &cum_offsets,
+  const paddle::Tensor &block_table,
+  int max_seq_len,
+  int max_dec_len,
+ float softmax_scale,
+  float in_scale,
+  bool causal,
+  cudaStream_t &stream,
+  paddle::Tensor *out);

csrc/gpu/append_attn/get_block_shape_and_split_kv_block.cu

@@ -14,6 +14,8 @@
 
 #include "helper.h"
 #include "paddle/extension.h"
+#include "utils.cuh"
+#include "cute/tensor.hpp"
 
 template <typename T>
 inline __device__ __host__ T div_up(T m, T n) {
@@ -48,6 +50,96 @@ __global__ void split_q_block(const int* __restrict__ seq_lens_q,
   }
 }
 
+template <uint32_t config_size>
+__global__ void search_chunk_size_for_mla(
+  const int * __restrict__ seq_lens_q,
+  const int * __restrict__ seq_lens_encoder,
+  const int * __restrict__ seq_lens_decoder,
+  int * __restrict__ num_blocks_x,
+  int * __restrict__ res_chunk_size,
+  const int bsz,
+  const int set_chunk_size,
+  const int block_size,
+  const int sm_cout) {
+  const uint32_t conf_id = threadIdx.x;
+  int gridx = 0;
+  if (set_chunk_size > 0 && conf_id == 0) {
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+      if (seq_len == 0 || seq_len_encoder > 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, set_chunk_size);
+      gridx += loop_times;
+    }
+    *num_blocks_x = gridx;
+    *res_chunk_size = set_chunk_size;
+  } else if (conf_id < config_size) {
+    __shared__ int gridx_shared[config_size];
+    // chunk_size is a multiple of 64
+    const int chunk_size = block_size << conf_id;
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+      if (seq_len == 0 || seq_len_encoder > 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
+      gridx += loop_times;
+    }
+    gridx_shared[conf_id] = gridx;
+    __syncthreads();
+    if (threadIdx.x == 0) {
+      uint32_t res_id = 0;
+      uint32_t max_last_wave_block = 0;
+      for (uint32_t i = 1; i < config_size; i++) {
+        uint32_t last_wave_block = gridx_shared[i] % sm_cout;
+        if (last_wave_block >= max_last_wave_block) {
+          res_id = i;
+          max_last_wave_block = last_wave_block;
+        }
+      }
+      *num_blocks_x = gridx_shared[res_id];
+      *res_chunk_size = block_size << res_id;
+    }
+  }
+}
+
+
+
+__global__ void split_block_for_mla(
+  const int * __restrict__ seq_lens_q,
+  const int * __restrict__ seq_lens_encoder,
+  const int * __restrict__ seq_lens_decoder,
+  int * __restrict__ batch_ids,
+  int * __restrict__ tile_ids_per_batch,
+  const int bsz,
+  const int chunk_size) {
+  if (threadIdx.x == 0) {
+    int index = 0;
+    for (uint32_t bid = 0; bid < bsz; bid++) {
+      int seq_len = seq_lens_q[bid];
+      int seq_len_encoder = seq_lens_encoder[bid];
+      int seq_len_decoder = seq_lens_decoder[bid] + seq_len;
+
+      if (seq_len == 0) continue;
+
+      int loop_times;
+      loop_times = cute::ceil_div(seq_len_decoder, chunk_size);
+      if (seq_len_encoder > 0) {
+        loop_times = 0;
+      }
+      for (uint32_t tile_id = 0; tile_id < loop_times; tile_id++) {
+        batch_ids[index] = bid;
+        tile_ids_per_batch[index++] = tile_id;
+      }
+    }
+  }
+}
+
 __global__ void split_kv_block(const int* __restrict__ seq_lens_decoder,
                                const int* __restrict__ seq_lens_encoder,
                                int* __restrict__ batch_ids,
@@ -110,7 +202,7 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
     const int decoder_step_token_num) {
   paddle::Tensor encoder_batch_ids, encoder_tile_ids_per_batch, encoder_num_blocks_x_cpu,
     kv_batch_ids, kv_tile_ids_per_batch, kv_num_blocks_x_cpu, decoder_batch_ids,
-    decoder_tile_ids_per_batch, decoder_num_blocks_x_cpu;
+    decoder_tile_ids_per_batch, decoder_num_blocks_x, decoder_num_blocks_x_cpu, decoder_chunk_size_device, decoder_chunk_size_cpu;
   auto stream = seq_lens_this_time.stream();
   int bsz = cum_offsets.shape()[0];
   const int encoder_block_shape_q = get_encoder_block_shape_q();
@@ -131,33 +223,91 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
   // decoder
   int max_dec_len_this_time_data = max_dec_len_this_time.data<int>()[0];
   if (max_dec_len_this_time_data > 0) {
-    const uint32_t decoder_max_tile_size_per_bs_q =
-        div_up((decoder_step_token_num * group_size), decoder_block_shape_q);
-    decoder_batch_ids =
-        GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q},
-                      paddle::DataType::INT32,
-                      seq_lens_encoder.place());
-    decoder_tile_ids_per_batch =
-        GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q},
-                      paddle::DataType::INT32,
-                      seq_lens_encoder.place());
-    auto decoder_num_blocks_x =
-        GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
-    split_q_block<<<1, 32, 0, stream>>>(seq_lens_this_time.data<int>(),
-                                        seq_lens_encoder.data<int>(),
-                                        decoder_batch_ids.data<int>(),
-                                        decoder_tile_ids_per_batch.data<int>(),
-                                        decoder_num_blocks_x.data<int>(),
-                                        bsz,
-                                        decoder_block_shape_q,
-                                        group_size);
-    decoder_num_blocks_x_cpu =
+    const bool mla_use_tensorcore = GetMlaUseTensorcore();
+    if (mla_use_tensorcore && group_size <=64) {
+      const int set_chunk_size = get_mla_dec_chunk_size(bsz);
+      decoder_chunk_size_device =
+          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+      decoder_num_blocks_x =
+          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+
+      int device;
+      cudaGetDevice(&device);
+      int sm_cout;
+      cudaDeviceGetAttribute(&sm_cout, cudaDevAttrMultiProcessorCount, device);
+      constexpr int config_size = 12; // search space for chunk size:[64, 128, 256, ... 131072]
+
+      search_chunk_size_for_mla<config_size><<<1, 32, 0, stream>>>(
+        seq_lens_this_time.data<int>(),
+        seq_lens_encoder.data<int>(),
+        seq_lens_decoder.data<int>(),
+        decoder_num_blocks_x.data<int>(),
+        decoder_chunk_size_device.data<int>(),
+        bsz,
+        set_chunk_size,
+        block_size,
+        sm_cout);
+
+      decoder_chunk_size_cpu =
+        decoder_chunk_size_device.copy_to(paddle::CPUPlace(), false);
+      decoder_num_blocks_x_cpu =
         decoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+
+      const int chunk_size = decoder_chunk_size_cpu.data<int>()[0];
+      const int num_blocks = decoder_num_blocks_x_cpu.data<int>()[0];
+      decoder_batch_ids = 
+          GetEmptyTensor({num_blocks},
+                        paddle::DataType::INT32,
+                        seq_lens_encoder.place());
+      decoder_tile_ids_per_batch =
+          GetEmptyTensor({num_blocks},
+                        paddle::DataType::INT32,
+                        seq_lens_encoder.place());
+
+      split_block_for_mla<<<1, 32, 0, stream>>>(
+        seq_lens_this_time.data<int>(),
+        seq_lens_encoder.data<int>(),
+        seq_lens_decoder.data<int>(),
+        decoder_batch_ids.data<int>(),
+        decoder_tile_ids_per_batch.data<int>(),
+        bsz,
+        chunk_size
+      );
+    } else {
+      const uint32_t decoder_max_tile_size_per_bs_q =
+          div_up((decoder_step_token_num * group_size), decoder_block_shape_q);
+      decoder_batch_ids =
+          GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q},
+                        paddle::DataType::INT32,
+                        seq_lens_encoder.place());
+      decoder_tile_ids_per_batch =
+          GetEmptyTensor({bsz * decoder_max_tile_size_per_bs_q},
+                        paddle::DataType::INT32,
+                        seq_lens_encoder.place());
+      decoder_num_blocks_x =
+          GetEmptyTensor({1}, paddle::DataType::INT32, seq_lens_encoder.place());
+      split_q_block<<<1, 32, 0, stream>>>(seq_lens_this_time.data<int>(),
+                                          seq_lens_encoder.data<int>(),
+                                          decoder_batch_ids.data<int>(),
+                                          decoder_tile_ids_per_batch.data<int>(),
+                                          decoder_num_blocks_x.data<int>(),
+                                          bsz,
+                                          decoder_block_shape_q,
+                                          group_size);
+      decoder_num_blocks_x_cpu =
+        decoder_num_blocks_x.copy_to(paddle::CPUPlace(), false);
+      decoder_chunk_size_cpu =
+        paddle::full({1}, 131072, paddle::DataType::INT32, paddle::CPUPlace());
+    }
   } else {
     decoder_batch_ids =
         paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
     decoder_tile_ids_per_batch =
         paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
+    decoder_chunk_size_cpu =
+        paddle::full({1}, 131072, paddle::DataType::INT32, paddle::CPUPlace());
+    decoder_num_blocks_x = 
+        paddle::full({1}, -1, paddle::DataType::INT32, paddle::GPUPlace());
     decoder_num_blocks_x_cpu =
         paddle::full({1}, -1, paddle::DataType::INT32, paddle::CPUPlace());
   }
@@ -230,7 +380,9 @@ std::vector<paddle::Tensor> GetBlockShapeAndSplitKVBlock(
           kv_num_blocks_x_cpu, /*cpu*/
           decoder_batch_ids,
           decoder_tile_ids_per_batch,
+          decoder_num_blocks_x,
           decoder_num_blocks_x_cpu, /*cpu*/
+          decoder_chunk_size_cpu, /*cpu*/
           max_len_kv_cpu /*cpu*/};
 }
 
@@ -250,6 +402,8 @@ std::vector<paddle::DataType> GetBlockShapeAndSplitKVBlockInferDtype(
           paddle::DataType::INT32,
           paddle::DataType::INT32,
           paddle::DataType::INT32,
+          paddle::DataType::INT32,
+          paddle::DataType::INT32,
           paddle::DataType::INT32};
 }
 
@@ -271,6 +425,8 @@ std::vector<std::vector<int64_t>> GetBlockShapeAndSplitKVBlockInferShape(
           dynamic_shape,
           dynamic_shape,
           {1},
+          {1},
+          {1},
           {1}};
 }
 
@@ -290,6 +446,8 @@ PD_BUILD_OP(get_block_shape_and_split_kv_block)
               "decoder_batch_ids",
               "decoder_tile_ids_per_batch",
               "decoder_num_blocks",
+              "decoder_num_blocks_cpu",
+              "decoder_chunk_size_cpu",
               "max_len_kv"})
     .Attrs({"group_size: int",
             "block_size: int",

csrc/gpu/append_attn/mla_cache_kernel.cu

@@ -134,6 +134,7 @@ std::vector<paddle::Tensor> DecodeMLAWriteCache(
                     const paddle::Tensor& cum_offsets,
                     const paddle::Tensor& block_tables,
                     const int max_seq_len,
+                    const bool speculate_decoder,
                     cudaStream_t& stream,
                     paddle::Tensor* kv_cache) {
   typedef PDTraits<T> traits_;
@@ -142,35 +143,59 @@ std::vector<paddle::Tensor> DecodeMLAWriteCache(
   
   auto max_blocks_per_seq = meta_data.max_blocks_per_seq;
   auto bsz = meta_data.batch_size;
+  auto token_num = meta_data.token_nums;
   auto block_size = meta_data.block_size;
   auto nope_size = meta_data.head_dims_v;
   auto all_size = meta_data.head_dims;
   int pe_size = all_size - nope_size;
   auto kv_num_heads = meta_data.kv_num_heads;
-  const uint32_t elem_nums = bsz * kv_num_heads * all_size;
-
   constexpr int PackSize = 16 / sizeof(DataType_);
-  const int pack_num = elem_nums / PackSize;
   const int blocksize = 128;
   int grid_size = 1;
-  GetNumBlocks<128>(pack_num, &grid_size);
 
-  decode_absorb_cache_kernel<DataType_, PackSize>
-      <<<grid_size, blocksize, 0, stream>>>(
-          reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
-          reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
-          reinterpret_cast<DataType_*>(kv_cache->data<data_t>()),
-          block_tables.data<int>(),
-          cum_offsets.data<int>(),
-          seq_lens.data<int>(),
-          seq_lens_encoder.data<int>(),
-          max_seq_len,
-          max_blocks_per_seq,
-          kv_num_heads,
-          nope_size,
-          pe_size,
-          block_size,
-          elem_nums);
+
+  if (speculate_decoder) {
+    const uint32_t elem_nums = token_num * kv_num_heads * all_size;
+    const int pack_num = elem_nums / PackSize;
+    GetNumBlocks<128>(pack_num, &grid_size);
+    speculate_decode_absorb_cache_kernel<DataType_, PackSize>
+        <<<grid_size, blocksize, 0, stream>>>(
+            reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
+            reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
+            reinterpret_cast<DataType_*>(kv_cache->data<data_t>()),
+            block_tables.data<int>(),
+            padding_offsets.data<int>(),
+            cum_offsets.data<int>(),
+            seq_lens.data<int>(),
+            seq_lens_encoder.data<int>(),
+            max_seq_len,
+            max_blocks_per_seq,
+            kv_num_heads,
+            nope_size,
+            pe_size,
+            block_size,
+            elem_nums);
+  } else {
+    const uint32_t elem_nums = bsz * kv_num_heads * all_size;
+    const int pack_num = elem_nums / PackSize;
+    GetNumBlocks<128>(pack_num, &grid_size);
+    decode_absorb_cache_kernel<DataType_, PackSize>
+        <<<grid_size, blocksize, 0, stream>>>(
+            reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_nope.data<data_t>())),
+            reinterpret_cast<DataType_*>(const_cast<data_t*>(kv_pe.data<data_t>())),
+            reinterpret_cast<DataType_*>(kv_cache->data<data_t>()),
+            block_tables.data<int>(),
+            cum_offsets.data<int>(),
+            seq_lens.data<int>(),
+            seq_lens_encoder.data<int>(),
+            max_seq_len,
+            max_blocks_per_seq,
+            kv_num_heads,
+            nope_size,
+            pe_size,
+            block_size,
+            elem_nums);
+  }
   return {};
 }
 
@@ -184,7 +209,8 @@ std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
     const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
     const std::string& cache_quant_type_str,
-    const int max_seq_len) {
+    const int max_seq_len,
+    const bool speculate_decoder) {
   cudaStream_t stream = kv_pe.stream();
   AppendAttnMetaData meta_data;
   const auto& kv_nope_dims = kv_nope.dims();
@@ -210,6 +236,7 @@ std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
                               cum_offsets,
                               block_tables,
                               max_seq_len,
+                              speculate_decoder,
                               stream,
                               const_cast<paddle::Tensor*>(&kv_cache));
     }
@@ -223,6 +250,7 @@ std::vector<paddle::Tensor> DecodeMLAWriteCacheKernel(
                               cum_offsets,
                               block_tables,
                               max_seq_len,
+                              speculate_decoder,
                               stream,
                               const_cast<paddle::Tensor*>(&kv_cache));
     }
@@ -258,5 +286,6 @@ PD_BUILD_OP(decode_mla_write_cache)
     .Outputs({"kv_cache_out"})
     .SetInplaceMap({{"kv_cache", "kv_cache_out"}})
     .Attrs({"cache_quant_type_str: std::string",
-            "max_seq_len: int"})
+            "max_seq_len: int",
+            "speculate_decoder: bool"})
     .SetKernelFn(PD_KERNEL(DecodeMLAWriteCacheKernel));

csrc/gpu/append_attn/mla_cache_kernel.cuh

@@ -88,6 +88,89 @@ __global__ void decode_absorb_cache_kernel(
   }
 }
 
+template <typename T, int VecSize = 1>
+__global__ void speculate_decode_absorb_cache_kernel(
+    const T* __restrict__ kv_nope,  // [bsz, kv_num_heads, pe_size] 512
+    const T* __restrict__ kv_pe,  // [bsz, kv_num_heads, nope_size] 64
+    T* __restrict__ kv_cache,    // [num_blocks, kv_num_heads, block_size,
+                                  // nope_size]                             
+    const int* __restrict__ block_tables,     // [bsz, max_blocks_per_seq]
+    const int* __restrict__ padding_offsets,
+    const int* __restrict__ cum_offsets,
+    const int* __restrict__ seq_lens,          // [bsz]
+    const int* __restrict__ seq_lens_encoder,  // [bsz]
+    const int max_seq_len,
+    const int max_blocks_per_seq,
+    const int kv_num_heads,
+    const int nope_size,
+    const int pe_size,
+    const int block_size,
+    const uint32_t elem_cnt) {
+  using LoadT = AlignedVector<T, VecSize>;
+  constexpr int HalfVecSize = VecSize / 2;
+  LoadT src_vec;
+
+  int64_t global_thread_idx = blockDim.x * blockIdx.x + threadIdx.x;
+  const uint32_t nope_hidden_size = kv_num_heads * nope_size;
+  const uint32_t pe_hidden_size = kv_num_heads * pe_size;
+  const uint32_t all_size = nope_size + pe_size;
+  const int64_t hidden_size = nope_hidden_size + pe_hidden_size;
+
+  for (int32_t linear_index = global_thread_idx * VecSize,
+               step = gridDim.x * blockDim.x * VecSize;
+       linear_index < elem_cnt;
+       linear_index += step) {
+    const int token_id = linear_index / hidden_size;
+    const int ori_bi = (token_id + padding_offsets[token_id]) / max_seq_len;
+    if (seq_lens[ori_bi] == 0) continue;
+    const int bias = linear_index % hidden_size;
+    const int start_token_idx = ori_bi * max_seq_len - cum_offsets[ori_bi];
+    const int write_seq_id = 
+        seq_lens[ori_bi] + token_id - start_token_idx;
+    if (write_seq_id == 0) continue;
+
+    const int* block_table_now = nullptr;
+
+    block_table_now = block_tables + ori_bi * max_blocks_per_seq;
+    const int block_idx = block_table_now[write_seq_id / block_size];
+    const int block_offset = write_seq_id % block_size;
+    if (block_idx < 0) {
+      printf(
+          "Fatal Error!!!, block idx %d when write_seq_id is %d\n some key var "
+          "%d %d %d %d\n",
+          block_idx,
+          write_seq_id,
+          ori_bi,
+          seq_lens[ori_bi],
+          token_id,
+          cum_offsets[ori_bi]);
+    }
+    if (bias < nope_hidden_size) { // pe
+      const uint32_t inner_bias = bias;
+      const uint32_t hi = inner_bias / nope_size;
+      const uint32_t h_bias = inner_bias % nope_size;
+      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
+                             hi * block_size * all_size +
+                             block_offset * all_size + h_bias;
+      const uint32_t ori_idx =
+          token_id * nope_hidden_size + inner_bias;
+      Load<T, VecSize>(&kv_nope[ori_idx], &src_vec);
+      Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
+    } else {
+      const uint32_t inner_bias = bias - nope_hidden_size;
+      const uint32_t hi = inner_bias / pe_size;
+      const uint32_t h_bias = inner_bias % pe_size;
+      const uint32_t tgt_idx = block_idx * kv_num_heads * block_size * all_size +
+                             hi * block_size * all_size +
+                             block_offset * all_size + nope_size + h_bias;
+      const uint32_t ori_idx =
+          token_id * pe_hidden_size + inner_bias;
+      Load<T, VecSize>(&kv_pe[ori_idx], &src_vec);
+      Store<T, VecSize>(src_vec, &kv_cache[tgt_idx]);
+    }
+  }
+}
+
 template <typename T, int VecSize = 1>
 __global__ void prefill_absorb_cache_kernel(
     const T* __restrict__ kv_nope,  // [bsz, kv_num_heads, pe_size] 512

csrc/gpu/append_attn/multi_head_latent_attention_kernel.h

@@ -36,32 +36,3 @@ void DecodeMLAAttentionKernel(
     bool causal,
     cudaStream_t &stream,
     paddle::Tensor *out);
-
-inline uint32_t get_max_partition_size(int bsz) {
-    static const char* max_partition_size_env = std::getenv("FLAGS_cascade_attention_max_partition_size");
-    static const uint32_t max_partition_size =
-            max_partition_size_env == nullptr ? 0 : std::stoul(std::string(max_partition_size_env));
-    return (max_partition_size != 0 ? max_partition_size : (bsz == 1 ? 128 : 512));
-}
-
-
-inline uint32_t get_cascade_attention_deal_each_time() {
-    static const char* cascade_attention_deal_each_time_env = std::getenv("FLAGS_cascade_attention_deal_each_time");
-    static const uint32_t cascade_attention_deal_each_time =
-            cascade_attention_deal_each_time_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_deal_each_time_env));
-    return (cascade_attention_deal_each_time != 0 ? cascade_attention_deal_each_time : 32);
-}
-
-inline uint32_t get_cascade_attention_num_stages() {
-    static const char* cascade_attention_num_stages_env = std::getenv("FLAGS_cascade_attention_num_stages");
-    static const uint32_t cascade_attention_num_stages =
-            cascade_attention_num_stages_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_num_stages_env));
-    return cascade_attention_num_stages != 0 ? cascade_attention_num_stages : 2;
-}
-
-inline uint32_t get_cascade_attention_num_threads() {
-    static const char* cascade_attention_num_threads_env = std::getenv("FLAGS_cascade_attention_num_threads");
-    static const uint32_t cascade_attention_num_threads =
-            cascade_attention_num_threads_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_num_threads_env));
-    return cascade_attention_num_threads != 0 ? cascade_attention_num_threads : 128;
-}

csrc/gpu/append_attn/utils.cuh

@@ -279,6 +279,11 @@ __forceinline__ __host__ __device__ void vec_cast<nv_bfloat16, float>(
       __VA_ARGS__                                       \
       break;                                            \
     }                                                   \
+    case 192: {                                         \
+      constexpr size_t HEAD_DIM = 192;                  \
+      __VA_ARGS__                                       \
+      break;                                            \
+    }                                                   \
     default: {                                          \
       PD_THROW("not support the head_dim: ", head_dim); \
     }                                                   \

csrc/gpu/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_gelu_and_mul.h

@@ -117,7 +117,7 @@ class LeftGELUAndMul {
   CUTLASS_HOST_DEVICE
   FragmentOutput operator()(FragmentAccumulator const &lhs,
                             FragmentAccumulator const &rhs) const {
-    // Convert source to interal compute numeric type
+    // Convert source to internal compute numeric type
     NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
         accumulator_to_compute;
 

csrc/gpu/cutlass_kernels/fp8_gemm_fused/dual_gemm/thread/left_silu_and_mul.h

@@ -117,7 +117,7 @@ class LeftSiLUAndMul {
   CUTLASS_HOST_DEVICE
   FragmentOutput operator()(FragmentAccumulator const &lhs,
                             FragmentAccumulator const &rhs) const {
-    // Convert source to interal compute numeric type
+    // Convert source to internal compute numeric type
     NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round>
         accumulator_to_compute;
 

csrc/gpu/cutlass_kernels/fp8_gemm_fused/dual_gemm/threadblock/dual_mma_base.h

@@ -92,7 +92,7 @@ class DualMmaBase {
                               Shape::kN / WarpGemm::kN,
                               Shape::kK / WarpGemm::kK>;
 
-  /// Number of warp-level GEMM oeprations
+  /// Number of warp-level GEMM operations
   static int const kWarpGemmIterations =
       (WarpGemm::kK / Operator0::Policy::MmaShape::kK);
 

csrc/gpu/cutlass_kernels/fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h

@@ -20,5 +20,10 @@
 #include "fuse_gemm_gelu_template.h"
 
 #include "fuse_gemm_act_template_3x.h"
+#include "fuse_block_gemm_act_template_3x.h"
 
 bool fp8_fp8_gemm_scale_bias_act(GemmEpilogueAllParams params);
+
+bool fp8_fp8_gemm_ptr_scale_bias_act(GemmEpilogueAllParams params);
+
+bool fp8_fp8_block_gemm_scale_bias_act(GemmEpilogueAllParams params);

csrc/gpu/cutlass_kernels/fp8_gemm_fused/fuse_block_gemm_act_template_3x.h

@@ -0,0 +1,183 @@
+// 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 "fp8_common.h"
+
+#include "cutlass/cutlass.h"
+#include "cutlass/numeric_types.h"
+
+#include "cute/tensor.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+#include "cutlass/gemm/kernel/tile_scheduler_params.h"
+#include "cutlass/epilogue/dispatch_policy.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+
+#include "cutlass_kernels/gemm/dispatch_policy.hpp"
+#include "cutlass_kernels/gemm/collective/collective_builder.hpp"
+
+using namespace cute;
+
+template <
+  typename InputType = phi::dtype::float8_e4m3fn,
+  typename OutType = phi::dtype::float16,
+  bool hasbias = false,
+  template <class> typename Activation = cutlass::epilogue::thread::Identity,
+  typename TileShape = Shape<_128, _128, _128>,
+  typename ClusterShape = Shape<_1, _2, _1>,
+  typename KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<1>,
+  typename EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative,
+  typename TileSchedule = cutlass::gemm::PersistentScheduler,
+  typename SM = cutlass::arch::Sm90
+>
+bool dispatch_fuse_block_gemm_c3x(GemmEpilogueAllParams params){
+  using ElementA = typename std::conditional_t<std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
+                                                              cutlass::float_e4m3_t,
+                                                              cutlass::float_e5m2_t>;
+  using ElementB = ElementA;
+  using ElementD = typename std::conditional_t<std::is_same_v<OutType, phi::dtype::bfloat16>,
+                                  cutlass::bfloat16_t,
+                                  cutlass::half_t>;
+  using ElementC = std::conditional_t<
+        hasbias,
+        ElementD,
+        void>;
+  
+  constexpr int ScaleMsPerTile = size<0>(TileShape{});
+  constexpr int ScaleGranularityM = size<0>(TileShape{}) / ScaleMsPerTile;
+
+  static constexpr bool IsStreamK = cute::is_same_v<TileSchedule, cutlass::gemm::StreamKScheduler>;
+
+  using LayoutA = cutlass::layout::RowMajor;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  using LayoutC = cutlass::layout::RowMajor;
+  using LayoutD = cutlass::layout::RowMajor;
+
+  using ElementAccumulator = float;
+  using ElementCompute = float;
+  using ElementScalar = float;
+
+  // 16B alignment lets us use TMA
+  static constexpr int AlignmentA = 16 / sizeof(ElementA);
+  static constexpr int AlignmentB = 16 / sizeof(ElementB);
+  static constexpr int AlignmentC = hasbias ? 16 / sizeof(ElementC) : 8;
+  static constexpr int AlignmentD = 16 / sizeof(ElementD);
+
+  static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
+
+  using FusionOperation = cutlass::epilogue::fusion::LinCombEltAct<Activation, ElementD, ElementCompute, ElementC, ElementScalar, RoundStyle>;
+
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      SM, cutlass::arch::OpClassTensorOp,
+      TileShape,
+      ClusterShape,
+      cutlass::epilogue::collective::EpilogueTileAuto,
+      ElementAccumulator, ElementCompute,
+      ElementC, LayoutC, AlignmentC,
+      ElementD, LayoutD, AlignmentD,
+      EpilogueSchedule,
+      FusionOperation
+    >::CollectiveOp;
+
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      SM, cutlass::arch::OpClassTensorOp,
+      ElementA, LayoutA, AlignmentA,
+      ElementB, LayoutB, AlignmentB,
+      ElementAccumulator,
+      TileShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      KernelSchedule
+    >::CollectiveOp;
+
+  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+      Shape<int,int,int,int>,
+      CollectiveMainloop,
+      CollectiveEpilogue,
+      TileSchedule
+  >;
+
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+  using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape;
+
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  //
+  // Data members
+  //
+
+  /// Initialization
+  StrideA stride_A{params.lda, cute::Int<1>{}, params.M * params.lda};
+  StrideB stride_B{params.ldb, cute::Int<1>{}, params.N * params.ldb};
+  StrideC stride_C{0, cute::Int<1>{}, 0};
+  StrideD stride_D{params.ldd, cute::Int<1>{}, params.M * params.ldd};
+
+  auto a_ptr = reinterpret_cast<ElementA*>(const_cast<void*>(params.A));
+  auto a_scale_ptr = reinterpret_cast<float*>(const_cast<void*>(params.A_scale));
+  auto b_ptr = reinterpret_cast<ElementB*>(const_cast<void*>(params.B));
+  auto b_scale_ptr = reinterpret_cast<float*>(const_cast<void*>(params.B_scale));
+  auto c_ptr = reinterpret_cast<ElementC*>(const_cast<void*>(params.bias));
+  auto d_ptr = reinterpret_cast<ElementD*>(params.D);
+
+  ProblemShapeType problem_size = ProblemShapeType{params.M, params.N, params.K, params.batch_count};
+
+  typename Gemm::Arguments arguments{
+    cutlass::gemm::GemmUniversalMode::kGemm,
+    problem_size,
+    {a_ptr, stride_A, b_ptr, stride_B,
+    a_scale_ptr, b_scale_ptr},
+    {{params.scale}, // epilogue.thread
+      c_ptr, stride_C, d_ptr, stride_D}
+  };
+  if constexpr (hasbias){
+    arguments.epilogue.thread.beta = 1.0;
+  }
+
+  if constexpr (IsStreamK) {
+    using DecompositionMode = typename cutlass::gemm::kernel::detail::
+        PersistentTileSchedulerSm90StreamKParams::DecompositionMode;
+    using ReductionMode = typename cutlass::gemm::kernel::detail::
+        PersistentTileSchedulerSm90StreamKParams::ReductionMode;
+
+    arguments.scheduler.decomposition_mode = DecompositionMode::StreamK;
+    arguments.scheduler.reduction_mode = ReductionMode::Nondeterministic;
+  }
+  
+
+  Gemm gemm_op;
+
+  cutlass::Status status = gemm_op.can_implement(arguments);
+  if (status != cutlass::Status::kSuccess) {
+    std::cout << "Gemm::can_implement() failed. " << cutlassGetStatusString(status) << std::endl;
+    return false;
+  }
+  size_t workspace_size = Gemm::get_workspace_size(arguments);
+  phi::Allocator* allocator = paddle::GetAllocator(params.place);
+  auto workspace = allocator->Allocate(workspace_size);
+
+  status = gemm_op(arguments, workspace->ptr(), params.stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::cout << "Gemm::run() failed." << cutlassGetStatusString(status) << std::endl;
+    return false;
+  }
+  return true;
+}

csrc/gpu/cutlass_kernels/fp8_gemm_fused/fuse_dual_gemm_act_template_3x.h

@@ -110,12 +110,12 @@ bool dispatch_dual_gemm_act_sm90(DualGemmEpilogueAllParams params) {
         ptr_A = reinterpret_cast<ElementB const*>(params.B0);
         ptr_B0 = reinterpret_cast<ElementA const*>(params.A);
     }
-    StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(arg_m, params.K, 1));
-    StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(arg_n, params.K, 1));
+    StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(arg_m, params.K, params.batch_count));
+    StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(arg_n, params.K, params.batch_count));
     StrideC stride_C;
-    StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(arg_m, arg_n, 1));
+    StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(arg_m, arg_n, params.batch_count));
 
-    typename Gemm::Arguments arguments = {cutlass::gemm::GemmUniversalMode::kGemm, {arg_m, arg_n, params.K, 1},
+    typename Gemm::Arguments arguments = {cutlass::gemm::GemmUniversalMode::kGemm, {arg_m, arg_n, params.K, params.batch_count},
         {ptr_A, stride_A, ptr_B0, ptr_B1, stride_B, params.scale0, params.scale1},
         {{}, // epilogue.thread
             nullptr, stride_C, reinterpret_cast<ElementOutput*>(params.D), stride_D}};

csrc/gpu/cutlass_kernels/fp8_gemm_fused/fuse_gemm_act_template_3x.h

@@ -115,17 +115,17 @@ bool dispatch_fuse_gemm_act_sm90(GemmEpilogueAllParams params){
   //
 
   /// Initialization
-  StrideA stride_A{params.lda, cute::Int<1>{}, 0};
-  StrideB stride_B{params.ldb, cute::Int<1>{}, 0};
+  StrideA stride_A{params.lda, cute::Int<1>{}, params.M * params.lda};
+  StrideB stride_B{params.ldb, cute::Int<1>{}, params.N * params.ldb};
   StrideC stride_C{0, cute::Int<1>{}, 0};
-  StrideD stride_D{params.ldd, cute::Int<1>{}, 0};
+  StrideD stride_D{params.ldd, cute::Int<1>{}, params.ldd * params.M};
 
   auto a_ptr = reinterpret_cast<ElementA*>(const_cast<void*>(params.A));
   auto b_ptr = reinterpret_cast<ElementB*>(const_cast<void*>(params.B));
   auto c_ptr = reinterpret_cast<ElementC*>(const_cast<void*>(params.bias));
   auto d_ptr = reinterpret_cast<ElementD*>(params.D);
 
-  ProblemShapeType problem_size = ProblemShapeType{params.M, params.N, params.K, 1};
+  ProblemShapeType problem_size = ProblemShapeType{params.M, params.N, params.K, params.batch_count};
 
   typename Gemm::Arguments arguments{
     cutlass::gemm::GemmUniversalMode::kGemm,
@@ -157,3 +157,188 @@ bool dispatch_fuse_gemm_act_sm90(GemmEpilogueAllParams params){
   }
   return true;
 }
+
+
+template <
+  typename InputType = phi::dtype::float8_e4m3fn,
+  typename OutType = phi::dtype::float16,
+  bool hasbias = false,
+  typename TileShape = Shape<_128, _128, _128>,
+  typename ClusterShape = Shape<_2, _1, _1>,
+  typename KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum,
+  typename EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized,
+  typename SM = cutlass::arch::Sm90
+>
+bool dispatch_fp8_gemm_ptr_scale_sm90(GemmEpilogueAllParams params){
+  using ElementA = typename std::conditional_t<std::is_same_v<InputType, phi::dtype::float8_e4m3fn>,
+                                                              cutlass::float_e4m3_t,
+                                                              cutlass::float_e5m2_t>;
+  using ElementB = ElementA;
+  using ElementD = typename std::conditional_t<std::is_same_v<OutType, phi::dtype::bfloat16>,
+                                  cutlass::bfloat16_t,
+                                  cutlass::half_t>;
+  using ElementC = std::conditional_t<
+        hasbias,
+        ElementD,
+        void>;
+
+  using LayoutA = cutlass::layout::RowMajor;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  using LayoutC = cutlass::layout::RowMajor;
+  using LayoutD = cutlass::layout::RowMajor;
+
+  using ElementAccumulator = float;
+  using ElementCompute = float;
+  using ElementScalar = float;
+
+  // 16B alignment lets us use TMA
+  static constexpr int AlignmentA = 16 / sizeof(ElementA);
+  static constexpr int AlignmentB = 16 / sizeof(ElementB);
+  static constexpr int AlignmentC = hasbias ? 16 / sizeof(ElementC) : 8;
+  static constexpr int AlignmentD = 16 / sizeof(ElementD);
+
+  static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
+
+  using ScaledEpilogue =  // B_scale * A_scale * acc
+    cutlass::epilogue::fusion::Sm90EVT<cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementD, ElementCompute, RoundStyle>, // B_scale * A_scale * acc
+      cutlass::epilogue::fusion::Sm90ScalarBroadcast<ElementScalar>, // B_scale
+      cutlass::epilogue::fusion::Sm90EVT<cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementCompute, ElementCompute, RoundStyle>, // A_scale * acc
+        cutlass::epilogue::fusion::Sm90ScalarBroadcast<ElementScalar>, // A_scale
+        cutlass::epilogue::fusion::Sm90AccFetch // acc
+      >
+    >;
+
+  using ScaledEpilogueBias =  // B_scale * A_scale * acc + bias
+    cutlass::epilogue::fusion::Sm90EVT<cutlass::epilogue::fusion::Sm90Compute<
+      cutlass::homogeneous_multiply_add, ElementD, ElementC, RoundStyle>,
+      cutlass::epilogue::fusion::Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,_0>>, // beta
+      cutlass::epilogue::fusion::Sm90SrcFetch<ElementC>,
+      cutlass::epilogue::fusion::Sm90EVT<cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementC, ElementCompute, RoundStyle>, // B_scale * A_scale * acc
+        cutlass::epilogue::fusion::Sm90ScalarBroadcast<ElementScalar>, // B_scale
+        cutlass::epilogue::fusion::Sm90EVT<cutlass::epilogue::fusion::Sm90Compute<cutlass::multiplies, ElementCompute, ElementCompute, RoundStyle>, // A_scale * acc
+          cutlass::epilogue::fusion::Sm90ScalarBroadcast<ElementScalar>, // A_scale
+          cutlass::epilogue::fusion::Sm90AccFetch // acc
+        >
+      >
+    >;
+
+  using CustomEpilogue = std::conditional_t<
+        hasbias,
+        ScaledEpilogueBias,
+        ScaledEpilogue>;
+
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      SM, cutlass::arch::OpClassTensorOp,
+      TileShape,
+      ClusterShape,
+      cutlass::epilogue::collective::EpilogueTileAuto,
+      ElementAccumulator, ElementCompute,
+      ElementC, LayoutC, AlignmentC,
+      ElementD, LayoutD, AlignmentD,
+      EpilogueSchedule,
+      CustomEpilogue
+    >::CollectiveOp;
+
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      SM, cutlass::arch::OpClassTensorOp,
+      ElementA, LayoutA, AlignmentA,
+      ElementB, LayoutB, AlignmentB,
+      ElementAccumulator,
+      TileShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      KernelSchedule
+    >::CollectiveOp;
+
+  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+      Shape<int,int,int,int>,
+      CollectiveMainloop,
+      CollectiveEpilogue,
+      cutlass::gemm::PersistentScheduler
+  >;
+
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+  using ProblemShapeType = typename Gemm::GemmKernel::ProblemShape;
+
+  using StrideA = typename Gemm::GemmKernel::StrideA;
+  using StrideB = typename Gemm::GemmKernel::StrideB;
+  using StrideC = typename Gemm::GemmKernel::StrideC;
+  using StrideD = typename Gemm::GemmKernel::StrideD;
+
+  //
+  // Data members
+  //
+
+  /// Initialization
+  StrideA stride_A{params.lda, cute::Int<1>{}, params.M * params.lda};
+  StrideB stride_B{params.ldb, cute::Int<1>{}, params.N * params.ldb};
+  StrideC stride_C{0, cute::Int<1>{}, 0};
+  StrideD stride_D{params.ldd, cute::Int<1>{}, params.ldd * params.M};
+
+  auto a_ptr = reinterpret_cast<ElementA*>(const_cast<void*>(params.A));
+  auto b_ptr = reinterpret_cast<ElementB*>(const_cast<void*>(params.B));
+  auto a_scale_ptr = reinterpret_cast<float*>(const_cast<void*>(params.A_scale));
+  auto b_scale_ptr = reinterpret_cast<float*>(const_cast<void*>(params.B_scale));
+  auto c_ptr = reinterpret_cast<ElementC*>(const_cast<void*>(params.bias));
+  auto d_ptr = reinterpret_cast<ElementD*>(params.D);
+
+  ProblemShapeType problem_size = ProblemShapeType{params.M, params.N, params.K, params.batch_count};
+
+
+  typename Gemm::Arguments arguments{
+    cutlass::gemm::GemmUniversalMode::kGemm,
+    problem_size,
+    {a_ptr, stride_A, b_ptr, stride_B},
+    {}
+  };
+
+  if constexpr (hasbias){
+    arguments.epilogue = {
+      {
+        {{1}, {}, {{_0{}, _0{}, _0{}}}},
+        {},
+        {
+          {{}, {b_scale_ptr}, {{_0{}, _0{}, _0{}}}},
+          {
+            {{}, {a_scale_ptr}, {{_0{}, _0{}, _0{}}}},
+            {},
+            {} // multiplies
+          },
+          {} // multiplies 
+        },
+        {}
+      }, // epilogue.thread
+      c_ptr, stride_C, d_ptr, stride_D};
+  }else{
+     arguments.epilogue = {
+      {
+        {{}, {b_scale_ptr}, {{_0{}, _0{}, _0{}}}},
+        {
+          {{}, {a_scale_ptr}, {{_0{}, _0{}, _0{}}}},
+          {},
+          {} // multiplies
+        },
+        {} // multiplies 
+      }, // epilogue.thread
+      c_ptr, stride_C, d_ptr, stride_D};
+  }
+
+  Gemm gemm_op;
+
+  cutlass::Status status = gemm_op.can_implement(arguments);
+  if (status != cutlass::Status::kSuccess) {
+    std::cout << "Gemm::can_implement() failed. " << cutlassGetStatusString(status) << std::endl;
+    return false;
+  }
+  size_t workspace_size = Gemm::get_workspace_size(arguments);
+  phi::Allocator* allocator = paddle::GetAllocator(params.place);
+  auto workspace = allocator->Allocate(workspace_size);
+
+  status = gemm_op(arguments, workspace->ptr(), params.stream);
+  if (status != cutlass::Status::kSuccess) {
+    std::cout << "Gemm::run() failed." << cutlassGetStatusString(status) << std::endl;
+    return false;
+  }
+  return true;
+}

csrc/gpu/cutlass_kernels/gemm/collective/collective_builder.hpp

@@ -0,0 +1,138 @@
+// 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.
+
+// Adapt from
+// https://github.com/vllm-project/vllm/blob/v0.7.2/csrc/cutlass_extensions/gemm/collective/collective_builder.hpp
+
+#pragma once
+
+#include "cutlass/gemm/collective/builders/sm90_gmma_builder.inl"
+
+#include "cutlass_kernels/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp"
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::gemm::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// GMMA_TMA_WS_SS (BlockScaled Builders)
+template <
+  class ElementA,
+  class GmemLayoutATag,
+  int AlignmentA,
+  class ElementB,
+  class GmemLayoutBTag,
+  int AlignmentB,
+  class ElementAccumulator,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class StageCountType,
+  int ScaleGranularityM
+>
+struct CollectiveBuilder<
+    arch::Sm90,
+    arch::OpClassTensorOp,
+    ElementA,
+    GmemLayoutATag,
+    AlignmentA,
+    ElementB,
+    GmemLayoutBTag,
+    AlignmentB,
+    ElementAccumulator,
+    TileShape_MNK,
+    ClusterShape_MNK,
+    StageCountType,
+    KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>,
+    cute::enable_if_t<
+      not detail::is_use_rmem_A<ElementA, GmemLayoutATag, ElementB, GmemLayoutBTag>()>
+> {
+  using KernelScheduleType = KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>;
+
+  static_assert(is_static<TileShape_MNK>::value);
+  static_assert(is_static<ClusterShape_MNK>::value);
+#ifndef CUTLASS_SM90_COLLECTIVE_BUILDER_SUPPORTED
+  static_assert(cutlass::detail::dependent_false<ElementA>, "Unsupported Toolkit for SM90 Collective Builder\n");
+#endif
+  static_assert(detail::is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, detail::tma_alignment_bytes>(),
+                "Should meet TMA alignment requirement\n");
+
+  static constexpr bool IsArrayOfPointersGemm = (cute::is_any_of_v<KernelScheduleType,
+                                                                   KernelPtrArrayTmaWarpSpecializedCooperative,
+                                                                   KernelPtrArrayTmaWarpSpecializedPingpong>);
+  static constexpr bool IsFP8Input = detail::is_input_fp8<ElementA, ElementB>();
+  static_assert((!IsFP8Input || !IsArrayOfPointersGemm),
+                "KernelTmaWarpSpecializedCooperativeFP8BlockScaledAccum is only compatible with FP8 Blocked Scaled version right now.");
+
+  // For fp32 types, map to tf32 MMA value type
+  using ElementAMma = cute::conditional_t<cute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
+  using ElementBMma = cute::conditional_t<cute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
+
+  static constexpr cute::GMMA::Major GmmaMajorA = detail::gmma_ss_tag_to_major_A<ElementAMma, GmemLayoutATag>();
+  static constexpr cute::GMMA::Major GmmaMajorB = detail::gmma_ss_tag_to_major_B<ElementBMma, GmemLayoutBTag>();
+
+  static constexpr bool IsCooperative = cute::is_any_of_v<KernelScheduleType,
+                                                          KernelTmaWarpSpecializedCooperative,
+                                                          KernelPtrArrayTmaWarpSpecializedCooperative,
+                                                          KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<ScaleGranularityM>>;
+  using AtomLayoutMNK = cute::conditional_t<IsCooperative,
+      Layout<Shape<_2,_1,_1>>, Layout<Shape<_1,_1,_1>>>;
+
+  using TiledMma = decltype(cute::make_tiled_mma(cute::GMMA::ss_op_selector<
+      ElementAMma, ElementBMma, ElementAccumulator, TileShape_MNK, GmmaMajorA, GmmaMajorB>(), AtomLayoutMNK{}));
+
+  using GmemTiledCopyA = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<1>(ClusterShape_MNK{})));
+  using GmemTiledCopyB = decltype(detail::sm90_cluster_shape_to_tma_atom(shape<0>(ClusterShape_MNK{})));
+
+  using SmemLayoutAtomA = decltype(detail::ss_smem_selector<
+      GmmaMajorA, ElementAMma, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
+  using SmemLayoutAtomB = decltype(detail::ss_smem_selector<
+      GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
+
+  static constexpr size_t TensorMapStorage = IsArrayOfPointersGemm ? sizeof(cute::TmaDescriptor) * 2 /* for A and B */ : 0;
+  static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage);
+
+  static constexpr int PipelineStages = detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes - KernelSmemCarveout,
+      ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{});
+  using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<PipelineStages, ClusterShape_MNK, KernelScheduleType, ScaleGranularityM>;
+
+  using SmemCopyAtomA = void;
+  using SmemCopyAtomB = void;
+
+  using CollectiveOp = CollectiveMma<
+      DispatchPolicy,
+      TileShape_MNK,
+      ElementA,
+      TagToStrideA_t<GmemLayoutATag>,
+      ElementB,
+      TagToStrideB_t<GmemLayoutBTag>,
+      TiledMma,
+      GmemTiledCopyA,
+      SmemLayoutAtomA,
+      SmemCopyAtomA,
+      cute::identity,
+      GmemTiledCopyB,
+      SmemLayoutAtomB,
+      SmemCopyAtomB,
+      cute::identity
+    >;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace cutlass::gemm::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

csrc/gpu/cutlass_kernels/gemm/collective/fp8_accumulation.hpp

@@ -0,0 +1,196 @@
+// 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.
+
+// adapted from: https://github.com/soundOfDestiny/cutlass/blob/a4208aa6958864923505cade9c63eb2a6daf16e5/include/cutlass/gemm/collective/fp8_accumulation.hpp
+
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "cute/algorithm/clear.hpp"
+#include "cute/tensor.hpp"
+
+//////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////FP8 Accumulation///////////////////////////
+//////////////////////////////////////////////////////////////////////////////
+/// This class provides API to promote (add) or scale (multiply_add) the results
+/// from the tensor core accumulators to the main accumulators when the number 
+/// of MMAs reaches the max number of MMA interval specified by user, after that
+/// the tensor core accumulators are zeroed.
+//////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::gemm::collective {
+
+template <
+    class EngineAccum,
+    class LayoutAccum>
+struct GmmaFP8AccumulationWithScale {  
+  using TensorAccum = cute::Tensor<EngineAccum, LayoutAccum>;
+  using ElementAccumulator = typename EngineAccum::value_type;
+
+  static_assert(is_static<LayoutAccum>::value, "Accumulator Layout should be static");
+  static_assert(is_rmem<TensorAccum>::value , "Accumulator tensor must be rmem resident.");
+
+private:
+  TensorAccum& accum_;
+  TensorAccum accum_temp_;
+
+  uint32_t accum_promotion_interval_;         // defines the max num of executed MMAs after which accum should be promoted.
+  uint32_t mma_count_per_mainloop_iteration_; // num of MMAs per k_tile of mainloop
+  uint32_t mma_count_;                        // current executed MMAs
+  uint32_t reset_accum_flag_;                 // accum needs to be zeroed or not. 
+
+  // promote or `add` the partial accumulators to main accumulator (FADD).
+  CUTLASS_DEVICE
+  void promote_core() {
+    warpgroup_wait<0>();
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accum_); ++i) {
+      accum_(i) += accum_temp_(i);
+    }
+  }
+
+  // `multiply` scale the partial accumulators and `add` to main accumulator (FFMA).
+  template <
+    class EngineScale,
+    class LayoutScale>
+  CUTLASS_DEVICE
+  void scale_core(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+    using TensorScale = cute::Tensor<EngineScale, LayoutScale>;
+
+    static_assert(is_static<LayoutScale>::value, "Scale Layout should be static");
+    static_assert(is_rmem<TensorScale>::value , "Scale tensor must be rmem resident.");
+
+    static_assert(LayoutAccum{}.shape() == LayoutScale{}.shape(), "Accumulator and scale must have same shape.");
+
+    warpgroup_wait<0>();
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accum_); ++i) {
+      accum_(i) += accum_temp_(i) * scale(i);
+    }
+  }
+
+public:
+  CUTLASS_DEVICE
+  GmmaFP8AccumulationWithScale(
+      TensorAccum &accum,
+      uint32_t accum_promotion_interval,
+      uint32_t mma_count_per_mainloop_iteration)
+      : accum_(accum), 
+        accum_promotion_interval_(accum_promotion_interval),
+        mma_count_per_mainloop_iteration_(mma_count_per_mainloop_iteration),
+        mma_count_(0), 
+        reset_accum_flag_(0) 
+  {
+    accum_temp_ = cute::make_fragment_like(accum);
+  }
+
+  //
+  // Methods (Common)
+  //
+
+  CUTLASS_DEVICE 
+  TensorAccum& operator()() {
+    return accum_temp_;
+  }
+
+  /// prepare the MMA accumulators when initialization or zeroing is required.
+  CUTLASS_DEVICE
+  bool prepare_if_needed() { 
+    return reset_accum_flag_;
+  }
+
+  //
+  // Methods (for FADD version)
+  //
+
+  /// promote (add) the results from the MMA accumulators to main accumulator if needed.
+  CUTLASS_DEVICE
+  void promote_if_needed() {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
+    if (reset_accum_flag_) {
+      promote_core();
+      mma_count_ = 0;
+    }
+  }
+
+  /// promote (add) the residue results from the MMA accumulators to main accumulator if needed.
+  CUTLASS_DEVICE
+  void promote_residue_if_needed() {
+    if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
+      promote_core();
+    }
+  }
+
+  //
+  // Methods (for FFMA version)
+  //
+
+  /// scale (multiply_add) the results from the MMA accumulators to main accumulator if needed.
+  template <
+    class EngineScale,
+    class LayoutScale>
+  CUTLASS_DEVICE
+  void scale_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = __shfl_sync(0xffffffff, mma_count_ == accum_promotion_interval_, 0);
+    if (reset_accum_flag_) {
+      scale_core(scale);
+      mma_count_ = 0;
+    }
+  }
+
+  /// scale (multiply_add) the residue results from the MMA accumulators to main accumulator if needed.
+  template <
+    class EngineScale,
+    class LayoutScale>
+  CUTLASS_DEVICE
+  void scale_residue_if_needed(const cute::Tensor<EngineScale, LayoutScale> &scale) {
+    if (__shfl_sync(0xffffffff, mma_count_ > 0, 0)) {
+      scale_core(scale);
+    }
+  }
+};
+
+} // namespace cutlass::gemm::collective

csrc/gpu/cutlass_kernels/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

@@ -0,0 +1,744 @@
+// 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.
+
+// Adapt from https://github.com/vllm-project/vllm/blob/v0.7.2/csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+// Adapted (Heavily) from: https://github.com/soundOfDestiny/cutlass/blob/9d997ce0dea4c5fa1a617db6b7ff29aa9235822c/include/cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/trace.h"
+#include "cutlass/numeric_types.h"
+
+#include "cute/arch/cluster_sm90.hpp"
+#include "cute/arch/copy_sm80.hpp"
+#include "cute/arch/copy_sm90.hpp"
+#include "cute/algorithm/functional.hpp"
+#include "cute/atom/mma_atom.hpp"
+#include "cute/algorithm/gemm.hpp"
+#include "cute/tensor_predicate.hpp"
+#include "cute/numeric/arithmetic_tuple.hpp"
+
+#include "cutlass_kernels/gemm/dispatch_policy.hpp"
+#include "cutlass_kernels/gemm/collective/fp8_accumulation.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::gemm::collective {
+using namespace cute;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// WarpSpecialized Mainloop
+template <
+  int Stages,
+  class ClusterShape,
+  class KernelSchedule,
+  int ScaleGranularityM_,
+  class TileShape_,
+  class ElementA_,
+  class StrideA_,
+  class ElementB_,
+  class StrideB_,
+  class TiledMma_,
+  class GmemTiledCopyA_,
+  class SmemLayoutAtomA_,
+  class SmemCopyAtomA_,
+  class TransformA_,
+  class GmemTiledCopyB_,
+  class SmemLayoutAtomB_,
+  class SmemCopyAtomB_,
+  class TransformB_>
+struct CollectiveMma<
+    MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_>,
+    TileShape_,
+    ElementA_,
+    StrideA_,
+    ElementB_,
+    StrideB_,
+    TiledMma_,
+    GmemTiledCopyA_,
+    SmemLayoutAtomA_,
+    SmemCopyAtomA_,
+    TransformA_,
+    GmemTiledCopyB_,
+    SmemLayoutAtomB_,
+    SmemCopyAtomB_,
+    TransformB_>
+{
+  //
+  // Type Aliases
+  //
+  using DispatchPolicy = MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8<Stages, ClusterShape, KernelSchedule, ScaleGranularityM_>;
+  using TileShape = TileShape_;
+  using ElementA = ElementA_;
+  using StrideA = StrideA_;
+  using ElementB = ElementB_;
+  using StrideB = StrideB_;
+  using TiledMma = TiledMma_;
+  using ElementAccumulator = typename TiledMma::ValTypeC;
+  using ElementBlockScale = ElementAccumulator;
+  using GmemTiledCopyA = GmemTiledCopyA_;
+  using GmemTiledCopyB = GmemTiledCopyB_;
+  using SmemLayoutAtomA = SmemLayoutAtomA_;
+  using SmemLayoutAtomB = SmemLayoutAtomB_;
+  using SmemCopyAtomA = SmemCopyAtomA_;
+  using SmemCopyAtomB = SmemCopyAtomB_;
+  using TransformA = TransformA_;
+  using TransformB = TransformB_;
+  using ArchTag = typename DispatchPolicy::ArchTag;
+
+  using CtaShape_MNK = decltype(shape_div(TileShape{}, ClusterShape{}));
+  using MainloopPipeline = cutlass::PipelineTmaAsync<DispatchPolicy::Stages>;
+  using PipelineState = cutlass::PipelineState<DispatchPolicy::Stages>;
+  using PipelineParams = typename MainloopPipeline::Params;
+
+  // Two threads per CTA are producers (1 for operand tile and 32 for scales)
+  static constexpr int NumProducerThreadEvents = 33; 
+
+  static constexpr int ScaleGranularityM = ScaleGranularityM_ == 0 ? size<0>(TileShape{}) : ScaleGranularityM_;
+  static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
+
+  static_assert(cute::rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  static_assert(cute::rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  static_assert((size<0>(TileShape{}) % ScaleGranularityM) == 0, "FP8 scaling granularity must evenly divide tile shape along M.");
+
+  // Tile along modes in a way that maximizes the TMA box size.
+  using SmemLayoutA = decltype(tile_to_shape(
+      SmemLayoutAtomA{},
+      make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
+      cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideA>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
+  using SmemLayoutB = decltype(tile_to_shape(
+      SmemLayoutAtomB{},
+      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{}),
+      cute::conditional_t< ::cutlass::gemm::detail::is_major<0,StrideB>(), Step<_2,_1,_3>, Step<_1,_2,_3>>{}));
+  
+  // Block scaling gmem-to-smem copy atom 
+  using SmemBlockScalingCopyAtomA = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
+  using SmemBlockScalingCopyAtomB = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<ElementBlockScale>, ElementBlockScale>;
+  
+  // Block scaling smem layout
+  using SmemLayoutScaleA = Layout<Shape<Int<ScaleMsPerTile>, Int<DispatchPolicy::Stages>>>;
+  using SmemLayoutScaleB = Layout<Shape<Int<DispatchPolicy::Stages>>, Stride<_1>>; // `ScaleNsPerTile` is always 1.
+
+  static_assert(DispatchPolicy::Stages >= 2, "Specialization requires Stages set to value 1 or more.");
+  static_assert(cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
+                cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
+                "MMA atom must source both A and B operand from smem_desc for this mainloop.");
+  static_assert(cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>,
+      "GmemTiledCopy - invalid SM90 TMA copy atom specified.");
+  static_assert(cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD> || cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>,
+      "GmemTiledCopy - invalid SM90 TMA copy atom specified.");
+  static_assert(cute::is_same_v<ElementAccumulator, ElementBlockScale>,
+             "ElementAccumulator and ElementBlockScale should be same datatype");
+
+  struct SharedStorage
+  {
+    struct TensorStorage : cute::aligned_struct<128> {
+      cute::array_aligned<typename TiledMma::ValTypeA, cute::cosize_v<SmemLayoutA>> smem_A;  // mxk
+      cute::array_aligned<typename TiledMma::ValTypeB, cute::cosize_v<SmemLayoutB>> smem_B;  // nxk
+      cute::array_aligned<ElementBlockScale, cute::cosize_v<SmemLayoutScaleA>> smem_scale_A; // ScaleMsPerTile x k
+      cute::array_aligned<ElementBlockScale, cute::cosize_v<SmemLayoutScaleB>> smem_scale_B; // 1xk
+    } tensors;
+
+    using PipelineStorage = typename MainloopPipeline::SharedStorage;
+    PipelineStorage pipeline;
+  };
+  using TensorStorage = typename SharedStorage::TensorStorage;
+  using PipelineStorage = typename SharedStorage::PipelineStorage;
+
+  // Host side kernel arguments
+  struct Arguments {
+    ElementA const* ptr_A;
+    StrideA dA;
+    ElementB const* ptr_B;
+    StrideB dB;
+    ElementBlockScale const* ptr_scale_A; 
+    ElementBlockScale const* ptr_scale_B;
+  };
+
+  // Device side kernel params
+  struct Params {
+    // Assumption: StrideA is congruent with Problem_MK
+    using TMA_A = decltype(make_tma_copy_A_sm90(
+        GmemTiledCopyA{},
+        make_tensor(static_cast<ElementA const*>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
+        SmemLayoutA{}(_,_,0),
+        TileShape{},
+        ClusterShape{}));
+    // Assumption: StrideB is congruent with Problem_NK
+    using TMA_B = decltype(make_tma_copy_B_sm90(
+        GmemTiledCopyB{},
+        make_tensor(static_cast<ElementB const*>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
+        SmemLayoutB{}(_,_,0),
+        TileShape{},
+        ClusterShape{}));
+    TMA_A tma_load_a;
+    TMA_B tma_load_b;
+    uint32_t tma_transaction_bytes = TmaTransactionBytes;
+    uint32_t tma_transaction_bytes_mk = TmaTransactionBytesMK;
+    uint32_t tma_transaction_bytes_nk = TmaTransactionBytesNK;
+    // Block scaling factors for A and B
+    ElementBlockScale const* ptr_scale_A; 
+    ElementBlockScale const* ptr_scale_B;
+  };
+
+  //
+  // Methods
+  //
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    (void) workspace;
+
+    // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    auto ptr_A = reinterpret_cast<ElementA const*>(args.ptr_A);
+    auto ptr_B = reinterpret_cast<ElementB const*>(args.ptr_B);
+
+    Tensor tensor_a = make_tensor(ptr_A, make_layout(make_shape(M,K,L), args.dA));
+    Tensor tensor_b = make_tensor(ptr_B, make_layout(make_shape(N,K,L), args.dB));
+    typename Params::TMA_A tma_load_a = make_tma_copy_A_sm90(
+        GmemTiledCopyA{},
+        tensor_a,
+        SmemLayoutA{}(_,_,cute::Int<0>{}),
+        TileShape{},
+        ClusterShape{});
+    typename Params::TMA_B tma_load_b = make_tma_copy_B_sm90(
+        GmemTiledCopyB{},
+        tensor_b,
+        SmemLayoutB{}(_,_,cute::Int<0>{}),
+        TileShape{},
+        ClusterShape{});
+    uint32_t transaction_bytes_mk = TmaTransactionBytesMK;
+    uint32_t transaction_bytes_nk = TmaTransactionBytesNK;
+    uint32_t transaction_bytes = transaction_bytes_mk + transaction_bytes_nk;
+
+    return {
+      tma_load_a,
+      tma_load_b,
+      transaction_bytes,
+      transaction_bytes_mk,
+      transaction_bytes_nk,
+      args.ptr_scale_A,
+      args.ptr_scale_B
+    };
+  }
+
+  template<class ProblemShape>
+  static bool
+  can_implement(
+      ProblemShape const& problem_shape,
+      [[maybe_unused]] Arguments const& args) {
+    constexpr int tma_alignment_bits = 128;
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    auto [M,N,K,L] = problem_shape_MNKL;
+    
+    bool implementable = true;
+    constexpr int min_tma_aligned_elements_A = tma_alignment_bits / cutlass::sizeof_bits<ElementA>::value;
+    implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_A>(cute::make_shape(M,K,L), StrideA{});
+    constexpr int min_tma_aligned_elements_B = tma_alignment_bits / cutlass::sizeof_bits<ElementB>::value;
+    implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_B>(cute::make_shape(N,K,L), StrideB{});
+
+    if (!implementable) {
+      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TMA.\n");
+    }
+    return implementable;
+  }
+
+  static constexpr int K_PIPE_MAX = DispatchPolicy::Stages;
+  static constexpr int K_PIPE_MMAS = 1;
+  static constexpr uint32_t TmaTransactionBytesMK =
+        cutlass::bits_to_bytes(size<0>(SmemLayoutA{}) * size<1>(SmemLayoutA{}) * static_cast<uint32_t>(sizeof_bits<ElementA>::value));
+  static constexpr uint32_t TmaTransactionBytesNK =
+        cutlass::bits_to_bytes(size<0>(SmemLayoutB{}) * size<1>(SmemLayoutB{}) * static_cast<uint32_t>(sizeof_bits<ElementB>::value));
+  static constexpr uint32_t TmaTransactionBytes = TmaTransactionBytesMK + TmaTransactionBytesNK;
+
+  /// Issue Tma Descriptor Prefetch -- ideally from a single thread for best performance
+  CUTLASS_DEVICE
+  static void prefetch_tma_descriptors(Params const& mainloop_params)
+  {
+    cute::prefetch_tma_descriptor(mainloop_params.tma_load_a.get_tma_descriptor());
+    cute::prefetch_tma_descriptor(mainloop_params.tma_load_b.get_tma_descriptor());
+  }
+
+  /// Set up the data needed by this collective for load and mma.
+  /// Returns a tuple of tensors. The collective and the kernel layer have the contract
+  /// Returned tuple must contain at least two elements, with the first two elements being:
+  /// gA_mkl - The tma tensor, A after a local tile so it has shape  (BLK_M,BLK_K,m,k,l)
+  /// gB_nkl - The tma tensor, B after a local tile so it has shape  (BLK_N,BLK_K,n,k,l)
+  template <class ProblemShape_MNKL>
+  CUTLASS_DEVICE auto
+  load_init(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params) const {
+    using X = Underscore;
+    // Separate out problem shape for convenience
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    // TMA requires special handling of strides to deal with coord codomain mapping
+    // Represent the full tensors -- get these from TMA
+    Tensor mA_mkl = mainloop_params.tma_load_a.get_tma_tensor(make_shape(M,K,L));                            // (m,k,l)
+    Tensor mB_nkl = mainloop_params.tma_load_b.get_tma_tensor(make_shape(N,K,L));                            // (n,k,l)
+
+    // Make tiled views, defer the slice
+    Tensor gA_mkl = local_tile(mA_mkl, TileShape{}, make_coord(_,_,_), Step<_1, X,_1>{});        // (BLK_M,BLK_K,m,k,l)
+    Tensor gB_nkl = local_tile(mB_nkl, TileShape{}, make_coord(_,_,_), Step< X,_1,_1>{});        // (BLK_N,BLK_K,n,k,l)
+
+    constexpr auto scales_m = Int<ScaleMsPerTile>{};
+    auto tM = get<2>(gA_mkl.shape());
+    auto tN = get<2>(gB_nkl.shape());
+    auto tK = get<3>(gA_mkl.shape());
+
+    // Make the tiled views of scale tensors
+    auto scaleA_shape = make_shape(M / ScaleGranularityM, tK, L); // (scale_m,k,l)
+    auto scaleA_layout = make_ordered_layout(scaleA_shape,  Step<_0, _1, _2>{});
+    auto scaleB_shape = make_shape(tN, tK, L); // (n,k,l)
+    auto scaleB_layout = make_ordered_layout(scaleB_shape, Step<_1, _0, _2>{});
+
+    // Note that mScaleA_mkl and mScaleB_nkl are already blocked tiled in the `m` host and 
+    // gScaleA_mkl and gScaleB_nkl in `g` global memory are same as mScaleA_mkl and mScaleB_nkl.
+    Tensor mScaleA_mkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_A), scaleA_layout); // (scale_m,k,l)
+    Tensor mScaleB_nkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_B), scaleB_layout); // (n,k,l)
+
+    return cute::make_tuple(gA_mkl, gB_nkl, mScaleA_mkl, mScaleB_nkl);
+  }
+
+  /// Perform a collective-scoped matrix multiply-accumulate
+  /// Producer Perspective
+  template <
+    class TensorA, class TensorB,
+    class TensorScaleA, class TensorScaleB,
+    class KTileIterator, class BlockCoord
+  >
+  CUTLASS_DEVICE void
+  load(
+      Params const& mainloop_params,
+      MainloopPipeline pipeline,
+      PipelineState smem_pipe_write,
+      cute::tuple<TensorA, TensorB, TensorScaleA, TensorScaleB> const& load_inputs,
+      BlockCoord const& blk_coord,
+      KTileIterator k_tile_iter, int k_tile_count,
+      int thread_idx,
+      uint32_t block_rank_in_cluster,
+      TensorStorage& shared_tensors) {
+    int lane_predicate = cute::elect_one_sync();
+
+    // Blockscaling: Tma loads for load_input and CpAsync for load_scale
+    Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{});        // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{});        // (BLK_N,BLK_K,PIPE)
+    Tensor sScaleA = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_A.data()), SmemLayoutScaleA{}); // (ScaleMsPerTile,k)
+    Tensor sScaleB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()), SmemLayoutScaleB{}); // (k)
+
+    //
+    // Prepare the TMA loads for A and B
+    //
+
+    constexpr uint32_t cluster_shape_x = get<0>(ClusterShape());
+    uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
+
+    Tensor gA_mkl = get<0>(load_inputs);
+    Tensor gB_nkl = get<1>(load_inputs);
+
+    auto block_tma_a = mainloop_params.tma_load_a.get_slice(cluster_local_block_id.y);
+    auto block_tma_b = mainloop_params.tma_load_b.get_slice(cluster_local_block_id.x);
+
+    // Partition the inputs based on the current block coordinates.
+    auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
+    Tensor gA = gA_mkl(_,_,m_coord,_,l_coord);                                                     // (BLK_M,BLK_K,k)
+    Tensor gB = gB_nkl(_,_,n_coord,_,l_coord);                                                     // (BLK_N,BLK_K,k)
+
+
+    // Block scaling: load_scale has scaling tensors in global memory which are not tiled
+    Tensor mScaleA_mkl = get<2>(load_inputs);
+    Tensor mScaleB_nkl = get<3>(load_inputs);
+    auto scales_m = get<0>(mScaleA_mkl.shape());
+
+    Tensor cScaleA_mkl = make_identity_tensor(mScaleA_mkl.shape());
+
+    Tensor gScaleA = local_tile( 
+      mScaleA_mkl, make_tile(Int<ScaleMsPerTile>{}), 
+      make_coord(m_coord,_,l_coord));                   // (ScaleMsPerTile,k,1)
+    Tensor cScaleA = local_tile( 
+      cScaleA_mkl, make_tile(Int<ScaleMsPerTile>{}), 
+      make_coord(m_coord,_,l_coord));
+    Tensor gScaleB = mScaleB_nkl(n_coord,_,l_coord);                                           // (1,k,1)
+
+    // TODO: test `scale_copy_a` with `ScaleMsPerTile` < 128
+    TiledCopy scale_copy_a = make_tiled_copy(SmemBlockScalingCopyAtomA{}, 
+      Layout<Shape<_32, _1>>{}, Layout<Shape<_4, _1>>{}); // (1,1,1)
+    TiledCopy scale_copy_b = make_tiled_copy(SmemBlockScalingCopyAtomB{}, 
+      Layout<Shape<_1>>{}, Layout<Shape<_1>>{}); // (1,1,1)
+    ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(threadIdx.x);
+    ThrCopy thr_scale_copy_b = scale_copy_b.get_slice(threadIdx.x);
+    
+    Tensor tAgA_ScaleA = thr_scale_copy_a.partition_S(gScaleA);
+    Tensor tAcA_ScaleA = thr_scale_copy_a.partition_S(cScaleA);
+    Tensor tAsA_ScaleA = thr_scale_copy_a.partition_D(sScaleA);
+    
+    Tensor tBgB_ScaleB = thr_scale_copy_b.partition_S(gScaleB);
+    Tensor tBsB_ScaleB = thr_scale_copy_b.partition_D(sScaleB);
+
+    // Applies the mapping from block_tma_a
+    Tensor tAgA = block_tma_a.partition_S(gA);                                              // (TMA,TMA_M,TMA_K,k)
+    Tensor tAsA = block_tma_a.partition_D(sA);                                              // (TMA,TMA_M,TMA_K,PIPE)
+
+    Tensor tBgB = block_tma_b.partition_S(gB);                                              // (TMA,TMA_N,TMA_K,k)
+    Tensor tBsB = block_tma_b.partition_D(sB);                                              // (TMA,TMA_N,TMA_K,PIPE)
+
+    uint16_t mcast_mask_a = 0;
+    uint16_t mcast_mask_b = 0;
+
+    // Issue TmaLoads for GEMM operands A/B and CpAsync for scale tensors
+    // Maps the tile -> block, value
+    if constexpr (cute::is_same_v<GmemTiledCopyA, SM90_TMA_LOAD_MULTICAST>) {
+      auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{};                       // (m,n) -> block_id
+      for (int n = 0; n < size<1>(block_layout); ++n) {
+        mcast_mask_a |= (uint16_t(1) << block_layout(cluster_local_block_id.x,n,Int<0>{}));
+      }
+    }
+
+    if constexpr (cute::is_same_v<GmemTiledCopyB, SM90_TMA_LOAD_MULTICAST>) {
+      auto block_layout = Layout<typename DispatchPolicy::ClusterShape>{};                       // (m,n) -> block_id
+      for (int m = 0; m < size<0>(block_layout); ++m) {
+        mcast_mask_b |= (uint16_t(1) << block_layout(m,cluster_local_block_id.y,Int<0>{}));
+      }
+    }
+
+    // Allocate predicate tensors for a_scales (since we can't guarantee that 
+    // all scales are valid, since we could have a partial tiles along M)
+    Tensor tApA_ScaleA = make_tensor<bool>(shape(tAsA_ScaleA(_,_,0)));
+    #pragma unroll
+    for (int i = 0; i < size(tApA_ScaleA); ++i) {
+      tApA_ScaleA(i) = get<0>(tAcA_ScaleA(i)) < scales_m;
+    }
+
+    // Mainloop
+    CUTLASS_PRAGMA_NO_UNROLL
+    for ( ; k_tile_count > 0; --k_tile_count) {
+      // LOCK smem_pipe_write for _writing_
+      pipeline.producer_acquire(smem_pipe_write);
+
+      //
+      // Copy gmem to smem for *k_tile_iter
+      //
+      int write_stage = smem_pipe_write.index();
+      using BarrierType = typename MainloopPipeline::ProducerBarrierType;
+      BarrierType* tma_barrier = pipeline.producer_get_barrier(smem_pipe_write);
+
+      // Copy operands A and B from global memory to shared memory
+      if (lane_predicate) copy(mainloop_params.tma_load_a.with(*tma_barrier, mcast_mask_a), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,write_stage));
+      if (lane_predicate) copy(mainloop_params.tma_load_b.with(*tma_barrier, mcast_mask_b), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,write_stage));
+
+      // Copy scale tensors from global memory to shared memory
+      copy_if(scale_copy_a, tApA_ScaleA, tAgA_ScaleA(_,_,*k_tile_iter), tAsA_ScaleA(_,_,write_stage));
+      copy(scale_copy_b, tBgB_ScaleB(_,*k_tile_iter), tBsB_ScaleB(_,write_stage));
+      pipeline.producer_commit(smem_pipe_write, cutlass::arch::cpasync_barrier_arrive_noinc);
+
+      ++k_tile_iter;
+
+      // Advance smem_pipe_write
+      ++smem_pipe_write;
+    }
+  }
+
+  /// Perform a Producer Epilogue to prevent early exit of blocks in a Cluster
+  CUTLASS_DEVICE void
+  load_tail(
+      MainloopPipeline pipeline,
+      PipelineState smem_pipe_write) {
+    int lane_predicate = cute::elect_one_sync();
+
+    // Issue the epilogue waits
+    if (lane_predicate) {
+      /* This helps avoid early exit of blocks in Cluster
+       * Waits for all stages to either be released (all
+       * Consumer UNLOCKs), or if the stage was never used
+       * then would just be acquired since the phase was
+       * still inverted from make_producer_start_state
+       */
+      pipeline.producer_tail(smem_pipe_write);
+    }
+  }
+
+  /// Perform a collective-scoped matrix multiply-accumulate
+  /// Consumer Perspective
+  template <
+    class FrgTensorC
+  >
+  CUTLASS_DEVICE void
+  mma(MainloopPipeline pipeline,
+      PipelineState smem_pipe_read,
+      FrgTensorC& accum,
+      int k_tile_count,
+      int thread_idx,
+      TensorStorage& shared_tensors,
+      Params const& mainloop_params) {
+
+
+    static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
+    static_assert(cute::rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
+    static_assert(cute::rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
+    static_assert(cute::is_void_v<SmemCopyAtomA>,
+      "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
+    static_assert(cute::is_void_v<SmemCopyAtomB>,
+      "SM90 GMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
+
+    Tensor sA = make_tensor(make_smem_ptr(shared_tensors.smem_A.data()), SmemLayoutA{});          // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(shared_tensors.smem_B.data()), SmemLayoutB{});          // (BLK_N,BLK_K,PIPE)
+    
+    // Block scaling
+    Tensor sScaleAViewAsC = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_A.data()),
+      Layout<
+        Shape<Shape<Int<ScaleGranularityM>, Int<ScaleMsPerTile>>, cute::tuple_element_t<1, TileShape>, Int<DispatchPolicy::Stages>>,
+        Stride<Stride<_0, _1>, _0, Int<ScaleMsPerTile>>
+      >{}); // ((ScaleGranularityM,ScaleMsPerTile),n,k)
+    Tensor sScaleB = make_tensor(cute::make_smem_ptr(shared_tensors.smem_scale_B.data()), SmemLayoutScaleB{}); // (k)
+
+    //
+    // Define C accumulators and A/B partitioning
+    //
+    
+    // Layout of warp group to thread mapping
+
+    static_assert(stride<0>(typename TiledMma::ALayout{}) == 0 and 
+                  stride<0>(typename TiledMma::BLayout{}) == 0 and
+                  size<0>(typename TiledMma::ALayout{}) == NumThreadsPerWarpGroup and
+                  size<0>(typename TiledMma::BLayout{}) == NumThreadsPerWarpGroup, 
+                  "Stride of the first mode must be 0 and the size of the mode must be NumThreadsPerWarpGroup");
+
+    constexpr int MmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
+    Layout warp_group_thread_layout = make_layout(Int<MmaWarpGroups>{}, 
+                                                  Int<NumThreadsPerWarpGroup>{});
+
+    int warp_group_idx = __shfl_sync(0xFFFFFFFF, thread_idx / NumThreadsPerWarpGroup, 0);
+
+    TiledMma tiled_mma;
+    auto thread_mma = tiled_mma.get_slice(warp_group_thread_layout(warp_group_idx));
+
+    Tensor tCsScaleAViewAsC = tiled_mma.get_slice(thread_idx).partition_C(sScaleAViewAsC);    // (MMA,MMA_M,MMA_N,PIPE), `thread_mma` above is correct when partitioning A and B, but it is not correct when partitioning C.
+
+    Tensor tCsA = thread_mma.partition_A(sA);                                                 // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCsB = thread_mma.partition_B(sB);                                                 // (MMA,MMA_N,MMA_K,PIPE)
+
+    // Allocate "fragments/descriptors"
+    Tensor tCrA = thread_mma.make_fragment_A(tCsA);                                           // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCrB = thread_mma.make_fragment_B(tCsB);                                           // (MMA,MMA_N,MMA_K,PIPE)
+
+    CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum));                                                         // M
+    CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum));                                                         // N
+    CUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));                                                          // K
+    CUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));                                                       // PIPE
+    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA));                                         // PIPE
+    CUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB));                                         // PIPE
+
+    //
+    // PIPELINED MAIN LOOP
+    //
+    static_assert((0 <= K_PIPE_MMAS) && (K_PIPE_MMAS <  K_PIPE_MAX),
+        "ERROR : Incorrect number of MMAs in flight");
+
+    // We release buffers to producer warps(dma load) with some mmas in flight
+    PipelineState smem_pipe_release = smem_pipe_read;
+    
+    // Per block scale values for operand A and B
+
+    using RegLayoutScaleAViewAsC = decltype(make_layout_like(tCsScaleAViewAsC(_, _, _, 0).layout())); // `make_layout_like` makes a compact layout.
+    using RegLayoutScaleAEssential = decltype(filter_zeros(RegLayoutScaleAViewAsC{}.stride(), RegLayoutScaleAViewAsC{}.shape())); // an interface to traverse the underlying storage for the compact layout mentioned above
+
+    Tensor tCrScaleAViewAsC = make_tensor<ElementBlockScale>(RegLayoutScaleAViewAsC{});              // (MMA,MMA_M,MMA_N)
+    ElementBlockScale scale_b;
+
+    // Prologue GMMAs
+    int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
+
+    tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
+
+    GmmaFP8AccumulationWithScale accumulation(accum, size<2>(TileShape{}) / size<2>(typename TiledMma::AtomShape_MNK{}), size<2>(tCrA));
+    warpgroup_fence_operand(accumulation());
+    CUTLASS_PRAGMA_UNROLL
+    for (int k_tile_prologue = prologue_mma_count; k_tile_prologue > 0; --k_tile_prologue)
+    {
+      // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
+      auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+      pipeline.consumer_wait(smem_pipe_read, barrier_token);
+
+      if (accumulation.prepare_if_needed()) {
+        tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
+      }
+
+      int read_stage = smem_pipe_read.index();
+      
+      // Load per block scale values from shared memory to registers.
+      scale_b = sScaleB[read_stage];
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
+        tCrScaleAViewAsC.data()[i] = tCsScaleAViewAsC(_, _, _, read_stage)(idx2crd(i, RegLayoutScaleAEssential{}));
+      }
+      if constexpr (ScaleMsPerTile == 1) {
+        static_assert(size(RegLayoutScaleAEssential{}) == 1);
+        tCrScaleAViewAsC.data()[0] = __shfl_sync(0xffffffff, tCrScaleAViewAsC.data()[0] * scale_b, 0); // `tCrScaleAViewAsC.data()[0]` are all same in a warp group when `ScaleMsPerTile == 1`.
+      } else {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
+          tCrScaleAViewAsC.data()[i] = tCrScaleAViewAsC.data()[i] * scale_b;
+        }
+      }
+
+      warpgroup_arrive();
+      // Unroll the K mode manually to set scale D to 1
+      CUTLASS_PRAGMA_UNROLL
+      for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+        // (V,M,K) x (V,N,K) => (V,M,N)
+        cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
+        tiled_mma.accumulate_ = GMMA::ScaleOut::One;
+      }
+      warpgroup_commit_batch();
+
+      // Block scale the accumulators with reg tensor `tCrScaleAViewAsC`
+      accumulation.scale_if_needed(tCrScaleAViewAsC);
+
+      ++smem_pipe_read;
+    }
+
+    warpgroup_fence_operand(accumulation());
+    // Mainloop GMMAs
+    k_tile_count -= prologue_mma_count;
+
+    CUTLASS_PRAGMA_NO_UNROLL
+    for ( ; k_tile_count > 0; --k_tile_count)
+    {
+      // WAIT on smem_pipe_read until its data are available (phase bit flips from rdPhaseBit value)
+      auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+      pipeline.consumer_wait(smem_pipe_read, barrier_token);
+
+      //
+      // Compute on k_tile
+      //
+
+      int read_stage = smem_pipe_read.index();
+
+      // Load per block scale values from shared memory to registers (at most twice per block along M and exactly once per block along N) 
+      scale_b = sScaleB[read_stage];
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
+        tCrScaleAViewAsC.data()[i] = tCsScaleAViewAsC(_, _, _, read_stage)(idx2crd(i, RegLayoutScaleAEssential{}));
+      }
+      if constexpr (ScaleMsPerTile == 1) {
+        static_assert(size(RegLayoutScaleAEssential{}) == 1);
+        tCrScaleAViewAsC.data()[0] = __shfl_sync(0xffffffff, tCrScaleAViewAsC.data()[0] * scale_b, 0); // `tCrScaleAViewAsC.data()[0]` are all same in a warp group when `ScaleMsPerTile == 1`.
+      } else {
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < size(RegLayoutScaleAEssential{}); i++) {
+          tCrScaleAViewAsC.data()[i] = tCrScaleAViewAsC.data()[i] * scale_b;
+        }
+      }
+
+      if (accumulation.prepare_if_needed()) {
+        tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
+      }
+
+      warpgroup_fence_operand(accumulation());
+      warpgroup_arrive();
+      // Unroll the K mode manually to set scale D to 1
+      CUTLASS_PRAGMA_UNROLL
+      for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+        // (V,M,K) x (V,N,K) => (V,M,N)
+        cute::gemm(tiled_mma, tCrA(_,_,k_block,read_stage), tCrB(_,_,k_block,read_stage), accumulation());
+        tiled_mma.accumulate_ = GMMA::ScaleOut::One;
+      }
+      warpgroup_commit_batch();
+
+      /// Wait on the GMMA barrier for K_PIPE_MMAS (or fewer) outstanding to ensure smem_pipe_write is consumed
+      warpgroup_wait<K_PIPE_MMAS>();
+      warpgroup_fence_operand(accumulation());
+
+      // Block scale the accumulators with reg tensor `tCrScaleAViewAsC`
+      accumulation.scale_if_needed(tCrScaleAViewAsC);
+
+      pipeline.consumer_release(smem_pipe_release);                 // UNLOCK smem_pipe_release, done _computing_ on it
+
+      // Advance smem_pipe_read and smem_pipe_release
+      ++smem_pipe_read;
+      ++smem_pipe_release;
+    }
+    
+    accumulation.scale_residue_if_needed(tCrScaleAViewAsC);
+
+    warpgroup_fence_operand(accumulation());
+  }
+
+  /// Perform a Consumer Epilogue to release all buffers
+  CUTLASS_DEVICE void
+  mma_tail(MainloopPipeline pipeline, PipelineState smem_pipe_release, int k_tile_count) {
+    // Prologue GMMAs
+    int prologue_mma_count = min(K_PIPE_MMAS, k_tile_count);
+    k_tile_count -= prologue_mma_count;
+
+    smem_pipe_release.advance(k_tile_count);
+
+    // Wait on all GMMAs to complete
+    warpgroup_wait<0>();
+
+    for (int count = 0; count < prologue_mma_count; ++count) {
+      pipeline.consumer_release(smem_pipe_release);                 // UNLOCK smem_pipe_release, done _computing_ on it
+      ++smem_pipe_release;
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass::gemm::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

csrc/gpu/cutlass_kernels/gemm/dispatch_policy.hpp

@@ -0,0 +1,53 @@
+// 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 "cutlass/gemm/dispatch_policy.hpp"
+
+namespace cutlass::gemm {
+
+//////////////////////////////////////////////////////////////////////////////
+
+// FP8 related policies (including Blocked Scaled Accumulation)
+//  `ScaleGranularityM` specifies scaling granularity along M, while zero-value
+//  `ScaleGranularityM` indicates that scaling granularity is
+//  `size<0>(TileShape_MNK{})` along M.
+template <int ScaleGranularityM = 0>
+struct KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum
+    : KernelTmaWarpSpecializedCooperative {};
+
+// n-buffer in smem (Hopper TMA), pipelined with Hopper GMMA and TMA, Warp
+// specialized dynamic schedule For FP8 kernels with Block Scaling
+template <int Stages_, class ClusterShape_ = Shape<_1, _1, _1>,
+          class KernelSchedule = KernelTmaWarpSpecialized,
+          int ScaleGranularityM =
+              0  // `ScaleGranularityM` specifies scaling granularity along M,
+                 // while zero-value `ScaleGranularityM` indicates that scaling
+                 // granularity is `size<0>(TileShape_MNK{})` along M.
+          >
+struct MainloopSm90TmaGmmaWarpSpecializedBlockScalingSubGroupMFP8
+    : MainloopSm90TmaGmmaWarpSpecialized<Stages_, ClusterShape_,
+                                         KernelSchedule> {
+  static_assert(
+      cute::is_same_v<
+          KernelSchedule,
+          KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<
+              ScaleGranularityM>>,
+      "KernelSchedule must be one of the warp specialized policies");
+};
+
+//////////////////////////////////////////////////////////////////////////////
+
+}  // namespace cutlass::gemm

csrc/gpu/cutlass_kernels/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_cooperative.hpp

@@ -291,7 +291,7 @@ public:
         args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN
             ? TileScheduler::RasterOrderOptions::AlongN
             : TileScheduler::RasterOrderOptions::AlongM;
-        return TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
+        return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
     }
 
     static dim3 get_block_shape()

csrc/gpu/cutlass_kernels/gemm/kernel/sm90_gemm_gated_tma_warpspecialized_pingpong.hpp

@@ -291,7 +291,7 @@ public:
         args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN
             ? TileScheduler::RasterOrderOptions::AlongN
             : TileScheduler::RasterOrderOptions::AlongM;
-        return TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
+        return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
     }
 
     static dim3 get_block_shape()

csrc/gpu/dequant_int8.cu

@@ -158,7 +158,7 @@ std::vector<paddle::DataType> DequantInt8Dtype(const paddle::DataType& input_dty
 }
 
 PD_BUILD_OP(dequant_int8)
-    .Inputs({"intput","out_scale"})
+    .Inputs({"input","out_scale"})
     .Outputs({"output"})
     .Attrs({"dtype: std::string"})
     .SetKernelFn(PD_KERNEL(DequantInt8))

csrc/gpu/env.h

@@ -0,0 +1,85 @@
+// 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
+
+inline uint32_t get_decoder_block_shape_q() {
+    static const char* decoder_block_shape_q_env = std::getenv("FLAGS_dec_block_shape_q");
+    static const uint32_t decoder_block_shape_q =
+            decoder_block_shape_q_env == nullptr ? 16 : std::stoi(std::string(decoder_block_shape_q_env));
+    return decoder_block_shape_q;
+}
+
+inline uint32_t get_encoder_block_shape_q() {
+    static const char* encoder_block_shape_q_env = std::getenv("FLAGS_enc_block_shape_q");
+    static const uint32_t encoder_block_shape_q =
+            encoder_block_shape_q_env == nullptr ? 64 : std::stoi(std::string(encoder_block_shape_q_env));
+    return encoder_block_shape_q;
+}
+
+inline uint32_t get_max_partition_size(int bsz) {
+    static const char* max_partition_size_env = std::getenv("FLAGS_cascade_attention_max_partition_size");
+    static const uint32_t max_partition_size =
+            max_partition_size_env == nullptr ? 32768 : std::stoul(std::string(max_partition_size_env));
+    return max_partition_size;
+}
+
+inline uint32_t get_cascade_attention_deal_each_time() {
+    static const char* cascade_attention_deal_each_time_env = std::getenv("FLAGS_cascade_attention_deal_each_time");
+    static const uint32_t cascade_attention_deal_each_time =
+            cascade_attention_deal_each_time_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_deal_each_time_env));
+    return (cascade_attention_deal_each_time != 0 ? cascade_attention_deal_each_time : 32);
+}
+
+inline uint32_t get_cascade_attention_num_stages() {
+    static const char* cascade_attention_num_stages_env = std::getenv("FLAGS_cascade_attention_num_stages");
+    static const uint32_t cascade_attention_num_stages =
+            cascade_attention_num_stages_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_num_stages_env));
+    return cascade_attention_num_stages != 0 ? cascade_attention_num_stages : 2;
+}
+
+inline uint32_t get_cascade_attention_num_threads() {
+    static const char* cascade_attention_num_threads_env = std::getenv("FLAGS_cascade_attention_num_threads");
+    static const uint32_t cascade_attention_num_threads =
+            cascade_attention_num_threads_env == nullptr ? 0 : std::stoul(std::string(cascade_attention_num_threads_env));
+    return cascade_attention_num_threads != 0 ? cascade_attention_num_threads : 128;
+}
+
+inline bool get_flags_mla_use_tensorcore() {
+    static const char* mla_use_tensorcore_env = std::getenv("FLAGS_mla_use_tensorcore");
+    static const uint32_t mla_use_tensorcore =
+            mla_use_tensorcore_env == nullptr ? 1 : std::stoul(std::string(mla_use_tensorcore_env));
+    return mla_use_tensorcore != 0 ? true : false;
+}
+
+inline bool get_mla_use_wg4() {
+    static const char* mla_use_wg4_env = std::getenv("FLAGS_mla_use_wg4");
+    static const uint32_t mla_use_wg4 =
+            mla_use_wg4_env == nullptr ? 1 : std::stoul(std::string(mla_use_wg4_env));
+    return mla_use_wg4 != 0 ? true : false;
+}
+
+inline bool enable_cuda_core_fp8_gemm() {
+    static const char* enable_cuda_core_fp8_env = std::getenv("FLAGS_cuda_core_fp8_gemm");
+    static const bool enable_cuda_core_fp8_gemm =
+            enable_cuda_core_fp8_env != nullptr && std::string(enable_cuda_core_fp8_env) == "1";
+    return enable_cuda_core_fp8_gemm;
+}
+
+inline int get_mla_dec_chunk_size(int bsz) {
+    static const char* mla_dec_chunk_size_env = std::getenv("FLAGS_mla_dec_chunk_size");
+    static const int mla_dec_chunk_size =
+            mla_dec_chunk_size_env == nullptr ? -1 : std::stoi(std::string(mla_dec_chunk_size_env));
+    return bsz > 1 ? mla_dec_chunk_size : 64;
+}

csrc/gpu/fp8_gemm_with_cutlass/fp8_common.h

@@ -47,6 +47,8 @@ typedef struct {
   std::vector<int64_t> &bias_dims;
   std::string &fuse_gemm_config;
   int split_k = 1;
+  const void *A_scale = nullptr;
+  const void *B_scale = nullptr;
 } GemmEpilogueAllParams;
 
 typedef bool (*func)(GemmEpilogueAllParams);

csrc/gpu/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu

@@ -0,0 +1,257 @@
+// 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 <iostream>
+
+#include "fp8_common.h"  // NOLINT
+#include "fp8_gemm_fused/fuse_block_gemm_act_template_3x.h"
+#include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
+
+std::vector<paddle::Tensor> cutlass_fp8_fp8_half_block_gemm_fused(
+    const paddle::Tensor& x,
+    const paddle::Tensor& y,
+    const paddle::Tensor& x_scale,
+    const paddle::Tensor& y_scale,
+    const paddle::optional<paddle::Tensor>& bias,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype,
+    std::string activation_type) {
+  paddle::Tensor out;
+  void* out_ptr = nullptr;
+  const void* x_ptr = nullptr;
+  const void* x_scale_ptr = nullptr;
+  const void* y_ptr = nullptr;
+  const void* y_scale_ptr = nullptr;
+
+
+  auto place = x.place();
+  cudaStream_t stream = x.stream();
+  int64_t device_id = place.GetDeviceId();
+  int sm_version = GetGPUComputeCapability(device_id);
+
+  int rank = x.dims().size();
+  int M = 0;
+  int K = 0;
+  int N = 0;
+  int ldd = 0;
+  
+  int lda = x.dims()[rank - 1];
+  int ldb = y.dims()[rank - 1];
+
+  if (!trans_x) {
+    M = x.dims()[rank - 2];
+    K = x.dims()[rank - 1];
+
+  } else {
+    M = x.dims()[rank - 1];
+    K = x.dims()[rank - 2];
+  }
+  if (!trans_y) {
+    N = y.dims()[rank - 1];
+    ldd = y.dims()[rank - 1];
+  } else {
+    N = y.dims()[rank - 2];
+    ldd = y.dims()[rank - 2];
+  }
+
+  int batch_count = 1;
+  for (size_t i = 0; i < rank - 2; ++i) {
+    batch_count *= x.dims()[i];
+  }
+
+  std::string input_dtype = "";
+  x_scale_ptr = reinterpret_cast<const void*>(x_scale.data<float>());
+  y_scale_ptr = reinterpret_cast<const void*>(y_scale.data<float>());
+  if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
+    input_dtype = "float8_e4m3fn";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
+    y_ptr = reinterpret_cast<const void*>(y.data<phi::dtype::float8_e4m3fn>());
+  } else if (x.dtype() == phi::DataType::FLOAT8_E5M2) {
+    input_dtype = "float8_e5m2";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e5m2>());
+    y_ptr = reinterpret_cast<const void*>(y.data<phi::dtype::float8_e5m2>());    
+  } else {
+    PADDLE_THROW(phi::errors::Fatal(
+        "fp8_fp8_half_gemm_fused only support e4m3 and e5m2 input"));
+  }
+
+  std::vector<int64_t> out_shape = x.shape();
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
+
+  if (output_dtype == "bfloat16") {
+    out = paddle::empty(out_shape, paddle::DataType::BFLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::bfloat16>());
+  } else if (output_dtype == "float16") {
+    out = paddle::empty(out_shape, paddle::DataType::FLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float16>());
+  } else {
+    PADDLE_THROW(phi::errors::Fatal(
+        "fp8_fp8_half_gemm_fused only support bfloat16 and float16 output"));
+  }
+
+  std::string isbias = bias ? "true" : "false";
+  std::string act = (activation_type == "" || activation_type == "identity")
+                        ? "noact"
+                        : activation_type;
+
+  std::string fuse_gemm_config =
+      input_dtype + "_" + output_dtype + "_" + isbias + "_" + act;
+
+  void* bias_data = nullptr;
+  std::vector<int64_t> bias_dims{};
+  if (bias) {
+    bias_dims = common::vectorize(bias.get().dims());
+    if (output_dtype == "bfloat16") {
+      bias_data = reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
+          bias.get().data<phi::dtype::bfloat16>()));
+    } else {
+      bias_data = reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
+          bias.get().data<phi::dtype::float16>()));
+    }
+  }
+
+    GemmEpilogueAllParams params = {x_ptr,
+                                    y_ptr,
+                                    out_ptr,
+                                    1,
+                                    M,
+                                    N,
+                                    K,
+                                    lda,
+                                    ldb,
+                                    ldd,
+                                    batch_count,
+                                    place,
+                                    stream,
+                                    sm_version,
+                                    0.01,  // for leaky_relu
+                                    bias_data,
+                                    bias_dims,
+                                    fuse_gemm_config,
+                                    0,
+                                    x_scale_ptr,
+                                    y_scale_ptr};
+    fp8_fp8_block_gemm_scale_bias_act(params);
+  // if(output_dtype == "bfloat16"){
+  //   dispatch_fuse_block_gemm_c3x<phi::dtype::float8_e4m3fn,
+  //                             phi::dtype::bfloat16,
+  //                             false,
+  //                             cutlass::epilogue::thread::Identity,
+  //                             cute::Shape<_128, _128, _128>,
+  //                             cute::Shape<_1, _2, _1>,
+  //                             cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8GroupBlockScaledAccum<1>,
+  //                             cutlass::epilogue::TmaWarpSpecializedCooperative,
+  //                             cutlass::arch::Sm90>(params);
+  // }else{
+  //   dispatch_fuse_block_gemm_c3x<phi::dtype::float8_e4m3fn,
+  //                             phi::dtype::float16,
+  //                             false,
+  //                             cutlass::epilogue::thread::Identity,
+  //                             cute::Shape<_128, _128, _128>,
+  //                             cute::Shape<_1, _2, _1>,
+  //                             cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8GroupBlockScaledAccum<1>,
+  //                             cutlass::epilogue::TmaWarpSpecializedCooperative,
+  //                             cutlass::arch::Sm90>(params);
+  // }
+  
+  return {out};
+}
+
+std::vector<std::vector<int64_t>> CutlassFp8Fp8HalfBlockGemmFusedInferShape(
+    const std::vector<int64_t>& x_shape,
+    const std::vector<int64_t>& y_shape,
+    const std::vector<int64_t>& x_scale_shape,
+    const std::vector<int64_t>& y_scale_shape,
+    const paddle::optional<std::vector<int64_t>>&  bias_shape,
+    bool trans_x,
+    bool trans_y){
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    y_shape.size(),
+                    phi::errors::InvalidArgument(
+                      "The rank of input X and Y should be equal, but received X's rank is %d, Y's rank is %d.",
+                      x_shape.size(),
+                      y_shape.size()));
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    x_scale_shape.size(),
+                    phi::errors::InvalidArgument(
+                      "The rank of input X and X_scale should be equal, but received X's rank is %d, X_scale's rank is %d.",
+                      x_shape.size(),
+                      x_scale_shape.size()));
+  PADDLE_ENFORCE_EQ(y_shape.size(),
+                    y_scale_shape.size(),
+                    phi::errors::InvalidArgument(
+                      "The rank of input Y and Y_scale should be equal, but received Y's rank is %d, Y_scale's rank is %d.",
+                      y_shape.size(),
+                      y_scale_shape.size()));
+  int rank = x_shape.size();
+  int M = 0;
+  int N = 0;
+  if ((x_shape[rank - 1] + 127) / 128 != x_scale_shape[rank - 2]){
+    PADDLE_THROW(phi::errors::Fatal(
+        "cutlass_fp8_fp8_half_block_gemm_fused only support x_scale's dim[-2] * 128 = x's dim[-1]."));
+  }
+  if (((y_shape[rank - 1] + 127) / 128 != y_scale_shape[rank - 1]) || ((y_shape[rank - 2] + 127) / 128 != y_scale_shape[rank - 2])){
+    PADDLE_THROW(phi::errors::Fatal(
+        "cutlass_fp8_fp8_half_block_gemm_fused only support input y_scale's dim[-2:] * 128 = y's dim[-2:]."));
+  }
+
+  if (!trans_x) {
+    M = x_shape[rank - 2];
+
+  } else {
+    M = x_shape[rank - 1];
+  }
+  if (!trans_y) {
+    N = y_shape[rank - 1];
+  } else {
+    N = y_shape[rank - 2];
+  }
+  std::vector<int64_t> out_shape = x_shape;
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
+  return {out_shape};
+}
+
+std::vector<paddle::DataType> CutlassFp8Fp8HalfBlockGemmFusedInferDtype(
+    const paddle::DataType& x_type,
+    const paddle::DataType& y_type,
+    const paddle::DataType& x_scale_type,
+    const paddle::DataType& y_scale_type,
+    const paddle::optional<paddle::DataType>& bias_type,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype) {
+    paddle::DataType data_type;
+    if (output_dtype == "bfloat16")
+        data_type = paddle::DataType::BFLOAT16;
+    else if (output_dtype ==  "float16")
+        data_type = paddle::DataType::FLOAT16;
+    else 
+        PD_THROW(
+                "cutlass_fp8_fp8_half_gemm only support bfloat16 and float16 output");
+    return {data_type};
+}
+
+PD_BUILD_OP(cutlass_fp8_fp8_half_block_gemm_fused)
+    .Inputs({"x", "y", "x_sacle", "y_scale", paddle::Optional("bias")})
+    .Attrs({"transpose_x: bool",
+            "transpose_y: bool",
+            "output_dtype: std::string",
+            "act: std::string"})
+    .Outputs({"out"})
+    .SetKernelFn(PD_KERNEL(cutlass_fp8_fp8_half_block_gemm_fused))
+    .SetInferShapeFn(PD_INFER_SHAPE(CutlassFp8Fp8HalfBlockGemmFusedInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(CutlassFp8Fp8HalfBlockGemmFusedInferDtype));

csrc/gpu/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.h

@@ -26,12 +26,5 @@ typedef struct {
     cudaStream_t stream;
 } GemmParams;
 
-inline bool enable_cuda_core_fp8_gemm() {
-    static const char* enable_cuda_core_fp8_env = std::getenv("FLAGS_cuda_core_fp8_gemm");
-    static const bool enable_cuda_core_fp8_gemm =
-            enable_cuda_core_fp8_env != nullptr && std::string(enable_cuda_core_fp8_env) == "1";
-    return enable_cuda_core_fp8_gemm;
-}
-
 template <typename InputType, typename OutputType>
 bool cuda_core_gemm_launcher(GemmParams const& params);

csrc/gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm_ptr_scale.cu

@@ -0,0 +1,213 @@
+// 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 <iostream>
+#include "fp8_common.h"
+
+#include "cutlass/cutlass.h"
+
+#include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
+
+
+std::vector<paddle::Tensor> cutlass_fp8_fp8_half_gemm_ptr_scale(
+    const paddle::Tensor& x,
+    const paddle::Tensor& y,
+    const paddle::optional<paddle::Tensor>& x_scale,
+    const paddle::optional<paddle::Tensor>& y_scale,
+    const paddle::optional<paddle::Tensor>& bias,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype) {
+  paddle::Tensor out;
+  void* out_ptr = nullptr;
+  const void* x_ptr = nullptr;
+  const void* x_scale_ptr = nullptr;
+  const void* y_ptr = nullptr;
+  const void* y_scale_ptr = nullptr;
+
+
+  auto place = x.place();
+  cudaStream_t stream = x.stream();
+  int64_t device_id = place.GetDeviceId();
+  int sm_version = GetGPUComputeCapability(device_id);
+
+  int rank = x.dims().size();
+  int M = 0;
+  int K = 0;
+  int N = 0;
+  int ldd = 0;
+  
+  int lda = x.dims()[rank - 1];
+  int ldb = y.dims()[rank - 1];
+
+  if (!trans_x) {
+    M = x.dims()[rank - 2];
+    K = x.dims()[rank - 1];
+
+  } else {
+    M = x.dims()[rank - 1];
+    K = x.dims()[rank - 2];
+  }
+  if (!trans_y) {
+    N = y.dims()[rank - 1];
+    ldd = y.dims()[rank - 1];
+  } else {
+    N = y.dims()[rank - 2];
+    ldd = y.dims()[rank - 2];
+  }
+
+  int batch_count = 1;
+  for (size_t i = 0; i < rank - 2; ++i) {
+    batch_count *= x.dims()[i];
+  }
+
+  std::string input_dtype = "";
+  x_scale_ptr = reinterpret_cast<const void*>(x_scale.get().data<float>());
+  y_scale_ptr = reinterpret_cast<const void*>(y_scale.get().data<float>());
+  if (x.dtype() == phi::DataType::FLOAT8_E4M3FN) {
+    input_dtype = "float8_e4m3fn";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e4m3fn>());
+    y_ptr = reinterpret_cast<const void*>(y.data<phi::dtype::float8_e4m3fn>());
+  } else if (x.dtype() == phi::DataType::FLOAT8_E5M2) {
+    input_dtype = "float8_e5m2";
+    x_ptr = reinterpret_cast<const void*>(x.data<phi::dtype::float8_e5m2>());
+    y_ptr = reinterpret_cast<const void*>(y.data<phi::dtype::float8_e5m2>());    
+  } else {
+    PADDLE_THROW(phi::errors::Fatal(
+        "fp8_fp8_half_gemm_fused only support e4m3 and e5m2 input"));
+  }
+
+  std::vector<int64_t> out_shape = x.shape();
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
+
+  if (output_dtype == "bfloat16") {
+    out = paddle::empty(out_shape, paddle::DataType::BFLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::bfloat16>());
+  } else if (output_dtype == "float16") {
+    out = paddle::empty(out_shape, paddle::DataType::FLOAT16, x.place());
+    out_ptr = reinterpret_cast<void*>(out.data<phi::dtype::float16>());
+  } else {
+    PADDLE_THROW(phi::errors::Fatal(
+        "fp8_fp8_half_gemm_fused only support bfloat16 and float16 output"));
+  }
+
+  std::string isbias = bias ? "true" : "false";
+  std::string act = "noact";
+
+  std::string fuse_gemm_config =
+      input_dtype + "_" + output_dtype + "_" + isbias + "_" + act;
+
+  void* bias_data = nullptr;
+  std::vector<int64_t> bias_dims{};
+  if (bias) {
+    bias_dims = common::vectorize(bias.get().dims());
+    if (output_dtype == "bfloat16") {
+      bias_data = reinterpret_cast<void*>(const_cast<phi::dtype::bfloat16*>(
+          bias.get().data<phi::dtype::bfloat16>()));
+    } else {
+      bias_data = reinterpret_cast<void*>(const_cast<phi::dtype::float16*>(
+          bias.get().data<phi::dtype::float16>()));
+    }
+  }
+
+  GemmEpilogueAllParams params = {x_ptr,
+                                    y_ptr,
+                                    out_ptr,
+                                    1,
+                                    M,
+                                    N,
+                                    K,
+                                    lda,
+                                    ldb,
+                                    ldd,
+                                    batch_count,
+                                    place,
+                                    stream,
+                                    sm_version,
+                                    0.01,  // for leaky_relu
+                                    bias_data,
+                                    bias_dims,
+                                    fuse_gemm_config,
+                                    0,
+                                    x_scale_ptr,
+                                    y_scale_ptr};
+  fp8_fp8_gemm_ptr_scale_bias_act(params);
+  return {out};
+}
+
+std::vector<std::vector<int64_t>> CutlassFp8Fp8HalfGemmPtrScaleFusedInferShape(
+    const std::vector<int64_t>& x_shape,
+    const std::vector<int64_t>& y_shape,
+    const paddle::optional<std::vector<int64_t>>&  x_scale_shape,
+    const paddle::optional<std::vector<int64_t>>&  y_scale_shape,
+    const paddle::optional<std::vector<int64_t>>&  bias_shape,
+    bool trans_x,
+    bool trans_y){
+  PADDLE_ENFORCE_EQ(x_shape.size(),
+                    y_shape.size(),
+                    phi::errors::InvalidArgument(
+                      "The rank of input X and Y should be equal, but received X's rank is %d, Y's rank is %d.",
+                      x_shape.size(),
+                      y_shape.size()));
+      
+  int rank = x_shape.size();
+  int M = 0;
+  int N = 0;
+
+  if (!trans_x) {
+    M = x_shape[rank - 2];
+  } else {
+    M = x_shape[rank - 1];
+  }
+  if (!trans_y) {
+    N = y_shape[rank - 1];
+  } else {
+    N = y_shape[rank - 2];
+  }
+  std::vector<int64_t> out_shape = x_shape;
+  out_shape[rank - 1] = N;
+  out_shape[rank - 2] = M;
+  return {out_shape};
+}
+
+std::vector<paddle::DataType> CutlassFp8Fp8HalfGemmPtrScaleFusedInferDtype(
+    const paddle::DataType& x_type,
+    const paddle::DataType& y_type,
+    const paddle::optional<paddle::DataType>& x_scale_type,
+    const paddle::optional<paddle::DataType>& y_scale_type,
+    const paddle::optional<paddle::DataType>& bias_type,
+    bool trans_x,
+    bool trans_y,
+    std::string output_dtype) {
+    paddle::DataType data_type;
+    if (output_dtype == "bfloat16")
+        data_type = paddle::DataType::BFLOAT16;
+    else if (output_dtype ==  "float16")
+        data_type = paddle::DataType::FLOAT16;
+    else 
+        PD_THROW(
+                "fp8_fp8_half_gemm_fused only support bfloat16 and float16 output");
+    return {data_type};
+}
+
+PD_BUILD_OP(cutlass_fp8_fp8_half_gemm_ptr_scale_fused)
+    .Inputs({"x", "y", paddle::Optional("x_scale"), paddle::Optional("y_scale"), paddle::Optional("bias")})
+    .Attrs({"transpose_x: bool",
+            "transpose_y: bool",
+            "output_dtype: std::string"})
+    .Outputs({"out"})
+    .SetKernelFn(PD_KERNEL(cutlass_fp8_fp8_half_gemm_ptr_scale))
+    .SetInferShapeFn(PD_INFER_SHAPE(CutlassFp8Fp8HalfGemmPtrScaleFusedInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(CutlassFp8Fp8HalfGemmPtrScaleFusedInferDtype));

csrc/gpu/fused_rotary_position_encoding.cu

@@ -101,6 +101,11 @@ void FusedRotaryPositionEncoding(
   int64_t query_stride = num_heads * head_size;
   int64_t key_stride = num_kv_heads * head_size;
 
+  if (num_tokens > 65535) {
+    PD_THROW(
+        "apply_rotary_embedding_kernel launch failed when num_tokens > 65535.");
+  }
+
   dim3 grid(num_tokens);
   dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
   PD_DISPATCH_FLOATING_AND_HALF_TYPES(

csrc/gpu/get_position_ids_and_mask_encoder_batch.cu

@@ -15,12 +15,13 @@
 #include "helper.h"
 #include "paddle/extension.h"
 
-__global__ void GetPositionIdsKernel(
+__global__ void GetPositionIdsAndMaskEncoderBatchKernel(
     const int* seq_lens_encoder,  // [bsz] 每个批次的 encoder 长度
     const int* seq_lens_decoder,  // [bsz] 每个批次的 decoder 长度
     const int* seq_lens_this_time,
-    int* position_ids,            // 输出的一维 position_ids
-    const int bsz) {              // 批次大小
+    int* position_ids,  // 输出的一维 position_ids
+    int* mask_encoder_batch,
+    const int bsz) {  // 批次大小
   // 当前线程索引（每个线程对应一个批次）
   int tid = threadIdx.x;
   if (tid >= bsz) return;
@@ -42,6 +43,7 @@ __global__ void GetPositionIdsKernel(
   // 写入 encoder 的 position_ids
   for (int i = 0; i < encoder_len; i++) {
     position_ids[offset + i] = i;
+    mask_encoder_batch[offset + i] = 1;
   }
   offset += encoder_len;
 
@@ -49,27 +51,36 @@ __global__ void GetPositionIdsKernel(
   if (decoder_len > 0) {
     for (int i = 0; i < seq_len_this_time; i++) {
       position_ids[offset + i] = decoder_len + i;  // 使用 decoder 长度本身
+      mask_encoder_batch[offset + i] = 0;
     }
   }
 }
 
 
-void GetPositionIds(const paddle::Tensor& seq_lens_encoder,
-                    const paddle::Tensor& seq_lens_decoder,
-                    const paddle::Tensor& seq_lens_this_time,
-                    const paddle::Tensor& position_ids) {
-  const int bsz = seq_lens_encoder.shape()[0];
+void GetPositionIdsAndMaskEncoderBatch(
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& position_ids,
+    const paddle::Tensor& mask_encoder_batch) {
+  const int bsz = seq_lens_this_time.shape()[0];
 
-  GetPositionIdsKernel<<<1, bsz, 0, position_ids.stream()>>>(
+  GetPositionIdsAndMaskEncoderBatchKernel<<<1, bsz, 0, position_ids.stream()>>>(
       seq_lens_encoder.data<int>(),
       seq_lens_decoder.data<int>(),
       seq_lens_this_time.data<int>(),
       const_cast<int*>(position_ids.data<int>()),
+      const_cast<int*>(mask_encoder_batch.data<int>()),
       bsz);
 }
 
-PD_BUILD_OP(get_position_ids)
-    .Inputs({"seq_lens_encoder", "seq_lens_decoder", "seq_lens_this_time", "position_ids"})
-    .Outputs({"position_ids_out"})
-    .SetInplaceMap({{"position_ids", "position_ids_out"}})
-    .SetKernelFn(PD_KERNEL(GetPositionIds));
+PD_BUILD_OP(get_position_ids_and_mask_encoder_batch)
+    .Inputs({"seq_lens_encoder",
+             "seq_lens_decoder",
+             "seq_lens_this_time",
+             "position_ids",
+             "mask_encoder_batch"})
+    .Outputs({"position_ids_out", "mask_encoder_batch_out"})
+    .SetInplaceMap({{"position_ids", "position_ids_out"},
+                    {"mask_encoder_batch", "mask_encoder_batch_out"}})
+    .SetKernelFn(PD_KERNEL(GetPositionIdsAndMaskEncoderBatch));

csrc/gpu/group_quant.cu

@@ -0,0 +1,201 @@
+// 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 "helper.h"
+#include<string.h>
+#include <cuda_runtime.h>
+
+template <typename T>
+__inline__ __device__ T WarpReduceAbsMax(T val, unsigned lane_mask) {
+  #pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1){
+    val = max(val, __shfl_xor_sync(lane_mask, val, mask, 32));
+  }
+  return val;
+}
+
+template <typename InType, typename OutType, int GroupSize, int VecSize, bool transpose_scale>
+__global__ void GroupQuantKernel(const InType* input,
+                                   const int64_t numel,
+                                   const int scale_rows,
+                                   const int scale_cols,
+                                   const float quant_max_bound,
+                                   const float quant_min_bound,
+                                   OutType* output,
+                                   float* out_scale_data) {
+
+  const int idx = (blockIdx.x * GroupSize) + threadIdx.x * VecSize;
+
+  if (idx >= numel) return;
+
+  const int scale_row_idx = blockIdx.x / scale_cols;
+  const int scale_col_idx = blockIdx.x % scale_cols;
+  int scale_idx = 0;
+  float abs_max_val = 0.000001;
+
+  if constexpr (transpose_scale) {
+    scale_idx = scale_col_idx * scale_rows + scale_row_idx;
+  }else{
+    scale_idx = scale_row_idx * scale_cols + scale_col_idx;
+  }
+
+  AlignedVector<InType, VecSize> in_vec;
+  AlignedVector<OutType, VecSize> out_vec;
+
+  Load<InType, VecSize>(input + idx, &in_vec);
+  #pragma unroll
+  for (int i = 0; i < VecSize; ++i) {
+    abs_max_val = max(abs_max_val, abs(static_cast<float>(in_vec[i])));
+  }
+  
+  abs_max_val = WarpReduceAbsMax(abs_max_val, 0xffffffff);
+  __syncthreads();
+  float scale = __shfl_sync(0xffffffff, abs_max_val, 0) / quant_max_bound;
+  
+  #pragma unroll
+  for (int i = 0; i < VecSize; ++i) {
+    float quant_val =  static_cast<float>(in_vec[i]) / scale;
+    quant_val = quant_val > quant_max_bound ? quant_max_bound : quant_val;
+    quant_val = quant_val < quant_min_bound ? quant_min_bound : quant_val;
+    out_vec[i] = static_cast<OutType>(quant_val);
+  }
+  Store<OutType, VecSize>(out_vec, output + idx);
+  if (threadIdx.x == 0) {
+    out_scale_data[scale_idx] = scale;
+  }
+}
+
+template <paddle::DataType InType, paddle::DataType OutType>
+std::vector<paddle::Tensor> LaunchGroupQuantKernel(const paddle::Tensor& x,
+                                                   const int group_size,
+                                                   const bool transpose_scale,
+                                                   const float quant_max_bound,
+                                                   const float quant_min_bound) {
+    typedef PDTraits<InType> in_traits;
+    typedef typename in_traits::DataType InDataType;
+    typedef typename in_traits::data_t in_data_t;
+
+    paddle::Tensor out;
+    paddle::Tensor scale_out;
+    auto place = x.place();
+    cudaStream_t stream = x.stream();
+    int rank = x.dims().size();
+    std::vector<int64_t> out_shape = x.shape();
+    std::vector<int64_t> scale_shape = x.shape();
+
+    out = paddle::empty(out_shape, OutType, place);
+    int64_t batch = 1;
+    int64_t m = x.shape()[rank - 2];
+    int64_t n = x.shape()[rank - 1];
+    int64_t scale_n = n / group_size;
+    for(auto i = 0; i < rank - 1; ++i){
+        batch *= x.shape()[i];
+    }
+    if(transpose_scale){
+        scale_shape[rank - 2] = scale_n;
+        scale_shape[rank - 1] = m;
+    }else{
+        scale_shape[rank - 1] = scale_n;
+    }
+    
+    scale_out = paddle::empty(scale_shape, paddle::DataType::FLOAT32, place);
+
+    int64_t numel = x.numel();
+    constexpr int VecSize = 4;
+
+    dim3 threadsPerBlock(group_size / VecSize, 1, 1);
+    dim3 block_per_grid((numel + group_size - 1) / group_size, 1, 1);
+
+    typedef PDTraits<OutType> out_traits;
+    typedef typename out_traits::DataType OutDataType;
+    typedef typename out_traits::data_t out_data_t;
+    
+    if(group_size == 128){
+        GroupQuantKernel<InDataType, OutDataType, 128, VecSize, true><<<block_per_grid, threadsPerBlock, 0, stream>>>(reinterpret_cast<const InDataType*>(x.data<in_data_t>()),
+                            numel,
+                            m,
+                            n / 128,
+                            quant_max_bound,
+                            quant_min_bound,
+                            reinterpret_cast<OutDataType*>(out.data<out_data_t>()),
+                            reinterpret_cast<float*>(scale_out.data<float>()));
+    }else{
+        PD_THROW("group_quant's group_size only support 128.");
+    }
+    
+    return {out, scale_out};
+}
+template <paddle::DataType InType>
+std::vector<paddle::Tensor> LaunchGroupQuant(const paddle::Tensor& x,
+                                             const int group_size,
+                                             const bool transpose_scale,
+                                             const float quant_max_bound,
+                                             const float quant_min_bound) {
+
+    if(fabs(quant_max_bound - 448.0f) < 0.000001){
+        return LaunchGroupQuantKernel<InType, paddle::DataType::FLOAT8_E4M3FN>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else{
+        PD_THROW("Only supported float8_e4m3fn quantization, please set quant_max_bound=448, quant_min_bound=-448.");
+    }
+    
+}
+
+
+std::vector<paddle::Tensor> GroupQuant(const paddle::Tensor& x,
+                                        const int group_size,
+                                        const bool transpose_scale,
+                                        const float quant_max_bound,
+                                        const float quant_min_bound) {
+    if(x.dtype() == paddle::DataType::FLOAT32){
+        return LaunchGroupQuant<paddle::DataType::FLOAT32>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else if(x.dtype() == paddle::DataType::FLOAT16){
+        return LaunchGroupQuant<paddle::DataType::FLOAT16>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else if(x.dtype() == paddle::DataType::BFLOAT16){
+        return LaunchGroupQuant<paddle::DataType::BFLOAT16>(x, group_size, transpose_scale, quant_max_bound, quant_min_bound);
+    }else{
+        PD_THROW("Unsupported data type.");
+    }
+}
+
+std::vector<std::vector<int64_t>> GroupQuantInferShape(const std::vector<int64_t>& input_shape, const int group_size, const bool transpose_scale, const float quant_max_bound,const float quant_min_bound) {
+    std::vector<int64_t> scale_shape = input_shape;
+    int rank = input_shape.size();
+    if(transpose_scale){
+        scale_shape[rank - 1] = input_shape[rank - 2];
+        scale_shape[rank - 2] = input_shape[rank - 1] / group_size;
+    }else{
+        scale_shape[rank - 1] = input_shape[rank - 1] / group_size;
+    }
+    return {input_shape, scale_shape};
+}
+
+std::vector<paddle::DataType> GroupQuantInferDtype(const paddle::DataType& input_dtype, const int group_size, const bool transpose_scale, const float quant_max_bound,const float quant_min_bound) {
+    
+    if(fabs(quant_max_bound - 448.0f) < 0.000001){
+        return {paddle::DataType::FLOAT8_E4M3FN, paddle::DataType::FLOAT32};
+    }else{
+        PD_THROW("Only supported attr of quant_max_bound in ['448.0'].");
+    }
+}
+
+PD_BUILD_OP(group_quant)
+    .Inputs({"x"})
+    .Outputs({"output", "scale"})
+    .Attrs({"group_size: int",
+            "transpose_scale: bool",
+            "quant_max_bound: float",
+            "quant_min_bound: float"})
+    .SetKernelFn(PD_KERNEL(GroupQuant))
+    .SetInferShapeFn(PD_INFER_SHAPE(GroupQuantInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(GroupQuantInferDtype));

csrc/gpu/helper.h

@@ -33,15 +33,19 @@ namespace cub = hipcub;
 #else
 #include <cub/cub.cuh>
 #include <curand_kernel.h>
+#include <cuda_fp8.h>
 #endif
 #include <iostream>
 #include <fstream>
 
+#include "env.h"
 #include "paddle/extension.h"
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/allocator.h"
+#include "paddle/phi/backends/gpu/gpu_info.h"
 #include "nlohmann/json.hpp"
 
+
 using json = nlohmann::json;
 
 #define CUDA_CHECK(call)                           \
@@ -151,6 +155,13 @@ public:
   typedef paddle::bfloat16 data_t;
 };
 
+template <>
+class PDTraits<paddle::DataType::FLOAT8_E4M3FN> {
+public:
+  typedef __nv_fp8_e4m3 DataType;
+  typedef paddle::float8_e4m3fn data_t;
+};
+
 template <typename T, int Size>
 struct alignas(sizeof(T) * Size) AlignedVector {
   T val[Size];
@@ -222,16 +233,15 @@ __device__ inline bool is_in_end(const int64_t id, const int64_t *end_ids, int l
     return flag;
 }
 
-inline uint32_t get_decoder_block_shape_q() {
-    static const char* decoder_block_shape_q_env = std::getenv("FLAGS_dec_block_shape_q");
-    static const uint32_t decoder_block_shape_q =
-            decoder_block_shape_q_env == nullptr ? 16 : std::stoi(std::string(decoder_block_shape_q_env));
-    return decoder_block_shape_q;
+inline int GetSMVersion() {
+  static int sm_version = phi::backends::gpu::GetGPUComputeCapability(
+      phi::backends::gpu::GetCurrentDeviceId());
+  return sm_version;
 }
 
-inline uint32_t get_encoder_block_shape_q() {
-    static const char* encoder_block_shape_q_env = std::getenv("FLAGS_enc_block_shape_q");
-    static const uint32_t encoder_block_shape_q =
-            encoder_block_shape_q_env == nullptr ? 64 : std::stoi(std::string(encoder_block_shape_q_env));
-    return encoder_block_shape_q;
+inline bool GetMlaUseTensorcore() {
+  static const bool flags_mla_use_tensorcore = get_flags_mla_use_tensorcore();
+  static const bool enable_mla_tensorcore = GetSMVersion() >= 90 ? true : false;
+  const bool mla_use_tensorcore = flags_mla_use_tensorcore && enable_mla_tensorcore;
+  return mla_use_tensorcore;
 }

csrc/gpu/int8_gemm_with_cutlass/epilogue_tensor_op_int32.h

@@ -176,7 +176,7 @@ struct DefaultIteratorsTensorOp<cutlass::bfloat16_t,
 ///
 /// Satisfies: ReadableTileIterator
 ///
-template <typename ThreadMap_  ///< Thread map (conept: OutputTileThreadMap)
+template <typename ThreadMap_  ///< Thread map (concept: OutputTileThreadMap)
           >
 class SharedLoadIteratorMixed<ThreadMap_, int32_t, 32, 16, 8, 8> {
  public:

csrc/gpu/mla_attn/attention_updater.cuh

@@ -0,0 +1,258 @@
+// 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.
+
+/*
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
+ * Dao. Licensed under the BSD 3-Clause.
+ *
+ * Modified by the FlashInfer team.
+ */
+
+#pragma once
+
+#include <cute/tensor.hpp>
+#include <cutlass/detail/helper_macros.hpp>
+
+#include "utils.cuh"
+
+namespace mla_attn {
+
+using namespace cute;
+
+template <typename T>
+struct MaxOp {
+  __device__ __forceinline__ T operator()(T const& x, T const& y) { return x > y ? x : y; }
+};
+
+template <>
+struct MaxOp<float> {
+  // This is slightly faster
+  __device__ __forceinline__ float operator()(float const& x, float const& y) { return max(x, y); }
+};
+
+template <typename T>
+struct SumOp {
+  __device__ __forceinline__ T operator()(T const& x, T const& y) { return x + y; }
+};
+
+template <int THREADS>
+struct Allreduce {
+  static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
+  template <typename T, typename Operator>
+  static __device__ __forceinline__ T run(T x, Operator& op) {
+    constexpr int OFFSET = THREADS / 2;
+    x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
+    return Allreduce<OFFSET>::run(x, op);
+  }
+};
+
+template <>
+struct Allreduce<2> {
+  template <typename T, typename Operator>
+  static __device__ __forceinline__ T run(T x, Operator& op) {
+    x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
+    return x;
+  }
+};
+
+template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
+          typename Operator>
+__device__ __forceinline__ void thread_reduce_(Tensor<Engine0, Layout0> const& tensor,
+                                               Tensor<Engine1, Layout1>& summary, Operator& op) {
+  static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+  static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+  CUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
+#pragma unroll
+  for (int mi = 0; mi < size<0>(tensor); mi++) {
+    summary(mi) = init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0));
+#pragma unroll
+    for (int ni = 1; ni < size<1>(tensor); ni++) {
+      summary(mi) = op(summary(mi), tensor(mi, ni));
+    }
+  }
+}
+
+template <typename Engine0, typename Layout0, typename Engine1, typename Layout1, typename Operator>
+__device__ __forceinline__ void quad_allreduce_(Tensor<Engine0, Layout0>& dst,
+                                                Tensor<Engine1, Layout1>& src, Operator& op) {
+  CUTE_STATIC_ASSERT_V(size(dst) == size(src));
+#pragma unroll
+  for (int i = 0; i < size(dst); i++) {
+    dst(i) = Allreduce<4>::run(src(i), op);
+  }
+}
+
+template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
+          typename Operator>
+__device__ __forceinline__ void reduce_(Tensor<Engine0, Layout0> const& tensor,
+                                        Tensor<Engine1, Layout1>& summary, Operator& op) {
+  thread_reduce_<init>(tensor, summary, op);
+  quad_allreduce_(summary, summary, op);
+}
+
+template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
+__device__ __forceinline__ void reduce_max(Tensor<Engine0, Layout0> const& tensor,
+                                           Tensor<Engine1, Layout1>& max) {
+  MaxOp<float> max_op;
+  reduce_<init>(tensor, max, max_op);
+}
+
+template <bool init, bool warp_reduce = true, typename Engine0, typename Layout0, typename Engine1,
+          typename Layout1>
+__device__ __forceinline__ void reduce_sum(Tensor<Engine0, Layout0> const& tensor,
+                                           Tensor<Engine1, Layout1>& sum) {
+  SumOp<float> sum_op;
+  thread_reduce_<init>(tensor, sum, sum_op);
+  if constexpr (warp_reduce) {
+    quad_allreduce_(sum, sum, sum_op);
+  }
+}
+
+template <typename Engine0, typename Layout0, typename Engine1, typename Layout1>
+__forceinline__ __device__ void apply_exp2(Tensor<Engine0, Layout0>& tensor,
+                                           Tensor<Engine1, Layout1> const& max) {
+  static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+  static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+  CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
+#pragma unroll
+  for (int mi = 0; mi < size<0>(tensor); ++mi) {
+    auto row_max = max(mi);
+#pragma unroll
+    for (int ni = 0; ni < size<1>(tensor); ++ni) {
+      tensor(mi, ni) = __expf(tensor(mi, ni) - row_max);
+    }
+  }
+}
+
+template <typename Engine0, typename Layout0, typename Engine1, typename Layout1>
+__forceinline__ __device__ void scale_apply_exp2(Tensor<Engine0, Layout0>& tensor,
+                                                 Tensor<Engine1, Layout1> const& max,
+                                                 const float scale) {
+  static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+  static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+  CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor));
+#pragma unroll
+  for (int mi = 0; mi < size<0>(tensor); ++mi) {
+    auto row_max = max(mi);
+#pragma unroll
+    for (int ni = 0; ni < size<1>(tensor); ++ni) {
+      // row_max * scale is a constant for each row, so we can use fma here
+      tensor(mi, ni) = __expf(tensor(mi, ni) * scale - row_max * scale);
+    }
+  }
+}
+
+template <int NUM_ROWS_PER_THREAD, bool WITH_SCALE>
+struct OnlineSoftmax {
+  constexpr static float fill_value = -5e4;
+  using TensorT = decltype(make_tensor<float>(Shape<Int<NUM_ROWS_PER_THREAD>>{}));
+  TensorT row_max, row_sum, scores_scale;
+  float sm_scale_log2;
+
+  CUTLASS_DEVICE OnlineSoftmax(float sm_scale_log2) : sm_scale_log2(sm_scale_log2) {
+    clear(scores_scale);
+  };
+
+  __forceinline__ __device__ TensorT get_lse() const { return row_sum; }
+
+  template <bool init, typename Tensor0>
+  __forceinline__ __device__ TensorT update(Tensor0& acc_s) {
+    // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))
+    Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+
+    static_assert(decltype(size<0>(scores))::value == NUM_ROWS_PER_THREAD);
+    if constexpr (init) {
+      reduce_max</*init=*/true>(scores, row_max);
+      if constexpr (WITH_SCALE) {
+        scale_apply_exp2(scores, row_max, sm_scale_log2);
+      } else {
+        apply_exp2(scores, row_max);
+      }
+      reduce_sum</*init=*/true, /*warp_reduce=*/false>(scores, row_sum);
+    } else {
+      // update row_max
+      Tensor scores_max_prev = make_fragment_like(row_max);
+      cute::copy(row_max, scores_max_prev);
+      reduce_max</*init=*/false>(scores, row_max);
+      // update scores_scale and scale row_sum
+#pragma unroll
+      for (int mi = 0; mi < size(row_max); ++mi) {
+        float scores_max_cur = row_max(mi);
+        if constexpr (WITH_SCALE) {
+          scores_scale(mi) = __expf((scores_max_prev(mi) - scores_max_cur) * sm_scale_log2);
+        } else {
+          scores_scale(mi) = __expf(scores_max_prev(mi) - scores_max_cur);
+        }
+        row_sum(mi) *= scores_scale(mi);
+      }
+      // perform exp2 on scores
+      if constexpr (WITH_SCALE) {
+        scale_apply_exp2(scores, row_max, sm_scale_log2);
+      } else {
+        apply_exp2(scores, row_max);
+      }
+      // update row_sum
+      reduce_sum</*init=*/false, /*warp_reduce=*/false>(scores, row_sum);
+      return scores_scale;
+    }
+  };
+
+  template <typename Tensor0>
+  __forceinline__ __device__ TensorT finalize(Tensor0& acc_s) {
+    // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))
+    Tensor scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+    static_assert(decltype(size<0>(scores))::value == NUM_ROWS_PER_THREAD);
+    SumOp<float> sum_op;
+    quad_allreduce_(row_sum, row_sum, sum_op);
+#pragma unroll
+    for (int mi = 0; mi < size(row_max); ++mi) {
+      float sum = row_sum(mi);
+      float inv_sum = 1.f / sum;
+      scores_scale(mi) = inv_sum;
+      row_max(mi) *= sm_scale_log2;
+    }
+    return scores_scale;
+  };
+
+  template <typename Tensor1>
+  __forceinline__ __device__ void rescale_o(Tensor1& acc_o) {
+    // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))
+    Tensor acc_o_rowcol = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+    static_assert(decltype(size<0>(acc_o_rowcol))::value == NUM_ROWS_PER_THREAD);
+#pragma unroll
+    for (int mi = 0; mi < size(row_max); ++mi) {
+#pragma unroll
+      for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
+        acc_o_rowcol(mi, ni) *= scores_scale(mi);
+      }
+    }
+  };
+
+  template <typename Tensor1, typename Tensor2>
+  __forceinline__ __device__ void rescale_o(Tensor1& acc_o, Tensor2& scores_scale_input) {
+    // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))
+    Tensor acc_o_rowcol = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+    static_assert(decltype(size<0>(acc_o_rowcol))::value == NUM_ROWS_PER_THREAD);
+#pragma unroll
+    for (int mi = 0; mi < size(row_max); ++mi) {
+#pragma unroll
+      for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) {
+        acc_o_rowcol(mi, ni) *= scores_scale_input(mi);
+      }
+    }
+  };
+};
+
+}  // namespace mla_attn
+

csrc/gpu/mla_attn/batch_mla_with_paged_kv_cache.cu

@@ -0,0 +1,298 @@
+// 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 <cuda.h>
+#include <cuda_device_runtime_api.h>
+#include <cutlass/arch/reg_reconfig.h>
+#include <cutlass/array.h>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/numeric_types.h>
+
+#include <type_traits>
+#include <vector>
+#include "cute/tensor.hpp"
+#include "mla_hopper.cuh"
+#include "mla_hopper_wg4.cuh"
+#include <iostream>
+#include <string>
+#include <sstream>
+
+#include "batch_mla_with_paged_kv_cache.h"
+#include "env.h"
+
+using namespace cute;
+using namespace mla_attn;
+using namespace std;
+
+template <typename T>
+struct cascade_type_traits {
+  using type = T;
+  using cutlass_type = T;
+};
+template <>
+struct cascade_type_traits<phi::dtype::bfloat16> {
+  using type = __nv_bfloat16;
+  using cutlass_type = cutlass::bfloat16_t;;
+};
+template <>
+struct cascade_type_traits<phi::dtype::float16> {
+  using type = half;
+  using cutlass_type = cutlass::half_t;
+};
+template <>
+struct cascade_type_traits<phi::dtype::float8_e4m3fn> {
+  using type = __nv_fp8_e4m3;
+  using cutlass_type = cutlass::float_e4m3_t;
+};
+
+template <typename T>
+void BatchMLAWithPagedKVCacheKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& q,  // [token_num, q_head_num, head_dim]
+    const paddle::Tensor& latent_cache,  // [max_block_num, q_head_num, block_size, head_dim]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const paddle::Tensor& num_blocks_x_device,
+    const std::string& cache_quant_type_str,
+    const int num_blocks_x,
+    const int chunk_size,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int draft_total_token_num,
+    const bool causal,
+    cudaStream_t& stream,
+    paddle::Tensor* out) {
+  using NV_TYPE = typename cascade_type_traits<T>::type;
+  using CUTLASS_TYPE = typename cascade_type_traits<T>::cutlass_type;
+  const auto token_num = meta_data.token_nums;
+  const auto block_size = meta_data.block_size;
+  const auto bsz = meta_data.batch_size;
+  const auto q_head_num = meta_data.q_num_heads;
+  const auto max_block_num_per_seq = meta_data.max_blocks_per_seq;
+  const auto max_block_num = bsz * max_block_num_per_seq;
+  const bool mla_use_wg4 = get_mla_use_wg4();
+#if CUDA_VERSION >= 12080
+  constexpr bool USE_REG_EALLOC = true;
+#else
+  constexpr bool USE_REG_EALLOC = false;
+#endif
+  int q_head_dim = meta_data.head_dims;
+  int k_head_dim = meta_data.head_dims;
+  int v_head_dim = meta_data.head_dims_v;
+  // int num_chunks = max_dec_len / chunk_size;
+  int num_chunks = div_up(max_dec_len, chunk_size);
+
+  auto *allocator = paddle::GetAllocator(q.place());
+  phi::Allocator::AllocationPtr O_tmp, m_tmp, d_tmp;
+  O_tmp = allocator->Allocate(
+      phi::SizeOf(q.dtype()) *
+      static_cast<size_t>(num_chunks * bsz * draft_total_token_num * q_head_num * v_head_dim));
+  m_tmp = allocator->Allocate(
+      sizeof(float) *
+      static_cast<size_t>(num_chunks * bsz * draft_total_token_num * q_head_num));
+  d_tmp = allocator->Allocate(
+      sizeof(float) *
+      static_cast<size_t>(num_chunks * bsz * draft_total_token_num * q_head_num));
+
+  if (q_head_dim == 576) {
+      if (mla_use_wg4 && block_size == 32) {
+        if constexpr (!std::is_same<CUTLASS_TYPE, cutlass::half_t>::value) {
+          PD_THROW("error!!! mla_wg4 not supports bf16 now!!!\n");
+        } else {
+          using ParamsType = Params<cutlass::half_t, cutlass::half_t, cutlass::half_t, int> ;
+          ParamsType params = {};
+          params.Q = reinterpret_cast<cutlass::half_t*>(const_cast<T*>(q.data<T>()));
+          params.KV = reinterpret_cast<cutlass::half_t*>(const_cast<T*>(latent_cache.data<T>()));
+          params.O = reinterpret_cast<cutlass::half_t*>(const_cast<T*>(out->data<T>()));
+          params.O_tmp = reinterpret_cast<cutlass::half_t*>(O_tmp->ptr());
+          params.m = reinterpret_cast<float*>(m_tmp->ptr());
+          params.d = reinterpret_cast<float*>(d_tmp->ptr());
+          params.block_tables = const_cast<int*>(block_tables.data<int>());
+          params.seq_lens_this_time = const_cast<int*>(seq_lens_this_time.data<int>());
+          params.seq_lens_encoder = const_cast<int*>(seq_lens_encoder.data<int>());
+          params.seq_lens_decoder = const_cast<int*>(seq_lens_decoder.data<int>());
+          params.cumsum_q_seqlens = const_cast<int*>(cu_seqlens_q.data<int>());
+          params.padding_offsets = const_cast<int*>(padding_offsets.data<int>());
+          params.batch_ids = const_cast<int*>(batch_ids.data<int>());
+          if (q_head_num <=64) {
+            params.tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+          } else {
+            params.q_tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+          }
+          params.num_blocks_x = const_cast<int*>(num_blocks_x_device.data<int>());
+          params.num_blocks_x_int = num_blocks_x;
+          params.q_stride_bsz = q_head_num * q_head_dim;
+          params.q_stride_head_num = q_head_dim;
+          params.kv_stride_block_num = block_size * k_head_dim;
+          params.kv_stride_block_size = k_head_dim;
+          params.o_stride_bsz = q_head_num * v_head_dim;
+          params.o_stride_head_num = v_head_dim;
+          params.bsz = bsz;
+          params.token_num = token_num;
+          params.max_seq_len = max_seq_len;
+          params.max_block_num = max_block_num;
+          params.max_block_num_per_seq = max_block_num_per_seq;
+          params.q_num_head = q_head_num;
+          params.qk_head_dim = q_head_dim;
+          params.vo_head_dim = v_head_dim;
+          params.block_size = block_size;
+          params.draft_total_token_num = draft_total_token_num;
+          params.sm_scale = softmax_scale;
+          params.chunk_size = chunk_size;
+          params.chunk_num = num_chunks;
+          BatchMLAWithPagedKVCacheWG4Dispatched<576, 512, half, ParamsType, USE_REG_EALLOC>(
+              params, stream
+          );
+        }
+      } else {
+        using ParamsType = Params<CUTLASS_TYPE, CUTLASS_TYPE, CUTLASS_TYPE, int>;
+        ParamsType params = {};
+        params.Q = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(q.data<T>()));
+        params.KV = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(latent_cache.data<T>()));
+        params.O = reinterpret_cast<CUTLASS_TYPE*>(const_cast<T*>(out->data<T>()));
+        params.O_tmp = reinterpret_cast<CUTLASS_TYPE*>(O_tmp->ptr());
+        params.m = reinterpret_cast<float*>(m_tmp->ptr());
+        params.d = reinterpret_cast<float*>(d_tmp->ptr());
+        params.block_tables = const_cast<int*>(block_tables.data<int>());
+        params.seq_lens_this_time = const_cast<int*>(seq_lens_this_time.data<int>());
+        params.seq_lens_encoder = const_cast<int*>(seq_lens_encoder.data<int>());
+        params.seq_lens_decoder = const_cast<int*>(seq_lens_decoder.data<int>());
+        params.cumsum_q_seqlens = const_cast<int*>(cu_seqlens_q.data<int>());
+        params.padding_offsets = const_cast<int*>(padding_offsets.data<int>());
+        params.batch_ids = const_cast<int*>(batch_ids.data<int>());
+        if (q_head_num <=64) {
+          params.tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+        } else {
+          params.q_tile_ids_per_batch = const_cast<int*>(tile_ids_per_batch.data<int>());
+        }
+        params.num_blocks_x = const_cast<int*>(num_blocks_x_device.data<int>());
+        params.num_blocks_x_int = num_blocks_x;
+        params.q_stride_bsz = q_head_num * q_head_dim;
+        params.q_stride_head_num = q_head_dim;
+        params.kv_stride_block_num = block_size * k_head_dim;
+        params.kv_stride_block_size = k_head_dim;
+        params.o_stride_bsz = q_head_num * v_head_dim;
+        params.o_stride_head_num = v_head_dim;
+        params.bsz = bsz;
+        params.token_num = token_num;
+        params.max_seq_len = max_seq_len;
+        params.max_block_num = max_block_num;
+        params.max_block_num_per_seq = max_block_num_per_seq;
+        params.q_num_head = q_head_num;
+        params.qk_head_dim = q_head_dim;
+        params.vo_head_dim = v_head_dim;
+        params.block_size = block_size;
+        params.draft_total_token_num = draft_total_token_num;
+        params.sm_scale = softmax_scale;
+        params.chunk_size = chunk_size;
+        params.chunk_num = num_chunks;
+        BatchMLAWithPagedKVCacheDispatched<576, 512, NV_TYPE, ParamsType, USE_REG_EALLOC>(
+            params, stream
+        );
+      }
+  } else {
+      PD_THROW("error!!! q_head_dim must be 576 !!!\n");
+  }
+}
+
+
+template void BatchMLAWithPagedKVCacheKernel<paddle::bfloat16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& q,  // [token_num, q_head_num, head_dim]
+    const paddle::Tensor& latent_cache,  // [max_block_num, q_head_num, block_size, head_dim]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const paddle::Tensor& num_blocks_x_device,
+    const std::string& cache_quant_type_str,
+    const int num_blocks_x,
+    const int chunk_size,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int draft_total_token_num,
+    const bool causal,
+    cudaStream_t& stream,
+    paddle::Tensor* out);
+
+
+template void BatchMLAWithPagedKVCacheKernel<paddle::float16>(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& q,  // [token_num, q_head_num, head_dim]
+    const paddle::Tensor& latent_cache,  // [max_block_num, q_head_num, block_size, head_dim]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const paddle::Tensor& num_blocks_x_device,
+    const std::string& cache_quant_type_str,
+    const int num_blocks_x,
+    const int chunk_size,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int draft_total_token_num,
+    const bool causal,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/mla_attn/batch_mla_with_paged_kv_cache.h

@@ -0,0 +1,70 @@
+// 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.
+
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * 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 "paddle/extension.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/allocator.h"
+#include "append_attn/utils.cuh"
+
+template <typename T>
+void BatchMLAWithPagedKVCacheKernel(
+    const AppendAttnMetaData& meta_data,
+    const paddle::Tensor& q,  // [token_num, q_head_num, head_dim]
+    const paddle::Tensor& latent_cache,  // [max_block_num, q_head_num, block_size, head_dim]
+    const paddle::optional<paddle::Tensor>& attn_mask,
+    const paddle::optional<paddle::Tensor>& cache_k_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_v_scale,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_k_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& cache_v_zp,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& shift_bias,  // [num_kv_heads, head_dim]
+    const paddle::optional<paddle::Tensor>& smooth_weight,  // [num_kv_heads, head_dim]
+    const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& seq_lens_decoder,
+    const paddle::Tensor& seq_lens_encoder,
+    const paddle::Tensor& cu_seqlens_q,
+    const paddle::Tensor& padding_offsets,
+    const paddle::Tensor& cum_offsets,
+    const paddle::Tensor& block_tables,
+    const paddle::Tensor& batch_ids,
+    const paddle::Tensor& tile_ids_per_batch,
+    const paddle::Tensor& num_blocks_x_device,
+    const std::string& cache_quant_type_str,
+    const int num_blocks_x,
+    const int chunk_size,
+    const int max_seq_len,
+    const int max_dec_len,
+    const float softmax_scale,
+    const float quant_max_bound,
+    const float quant_min_bound,
+    const float in_scale,
+    const int draft_total_token_num,
+    const bool causal,
+    cudaStream_t& stream,
+    paddle::Tensor* out);

csrc/gpu/mla_attn/epilogue.cuh

@@ -0,0 +1,178 @@
+// 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.
+
+/*
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
+ * Dao. Licensed under the BSD 3-Clause.
+ *
+ * Modified by the FlashInfer team.
+ */
+
+#pragma once
+
+#ifndef ATTENTION_HOPPER_EPILOGUE_CUH_
+#define ATTENTION_HOPPER_EPILOGUE_CUH_
+
+#include <cutlass/cutlass.h>
+
+#include "cute/tensor.hpp"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "named_barrier.cuh"
+#include "utils.cuh"
+
+#ifdef DEBUG_MLA
+#undef DEBUG_MLA
+#endif
+// #define DEBUG_MLA
+
+namespace mla_attn {
+
+using namespace cute;
+
+template <typename Ktraits>
+struct CollectiveEpilogue {
+  using DTypeO = typename Ktraits::DTypeO;
+  static constexpr int BLOCK_SHAPE_Q = Ktraits::BLOCK_SHAPE_Q;
+  static constexpr int BLOCK_SHAPE_KV = Ktraits::BLOCK_SHAPE_KV;
+  static constexpr int HEAD_DIM_VO = Ktraits::HEAD_DIM_VO;
+  using TileShape_PDV = Shape<Int<BLOCK_SHAPE_Q>, Int<HEAD_DIM_VO>, Int<BLOCK_SHAPE_KV>>;
+
+  static constexpr int NUM_WARPS = Ktraits::NUM_WARPS;
+  static constexpr int NUM_THREADS = NUM_WARPS * cutlass::NumThreadsPerWarp;
+
+  static constexpr int NUM_COPY_THREADS = cutlass::NumThreadsPerWarpGroup;
+  static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS;
+
+  using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_PDV{})),
+                                   decltype(cute::get<1>(TileShape_PDV{}))>());
+  using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{})));
+
+  using SmemCopyAtomO = Copy_Atom<cute::SM90_U32x4_STSM_N, DTypeO>;
+  using SharedStorage = cute::array_aligned<DTypeO, cute::cosize_v<SmemLayoutO>>;
+
+  using ShapeT = cute::Shape<int32_t, int32_t, int32_t>;
+  using StrideT = cute::Shape<int32_t, _1, int32_t>;
+  using LayoutT = cute::Layout<ShapeT, StrideT>;
+
+  using ShapeTmpT = cute::Shape<int32_t, int32_t, int32_t, int32_t>;
+  using StrideTmpT = cute::Shape<int32_t, _1, int32_t, int32_t>;
+  using LayoutTmpT = cute::Layout<ShapeTmpT, StrideTmpT>;
+
+  using ShapeNTMAT = cute::Shape<int32_t, int32_t>;
+  using StrideNTMAT = cute::Shape<int32_t, _1>;
+  using LayoutNTMAT = cute::Layout<ShapeNTMAT, StrideNTMAT>;
+
+  using GmemTiledCopyOTMA = cute::SM90_TMA_STORE;
+  using TMA_O = decltype(make_tma_copy(
+      GmemTiledCopyOTMA{},
+      make_tensor(make_gmem_ptr(static_cast<DTypeO*>(nullptr)), ShapeT{}, StrideT{}), SmemLayoutO{},
+      select<0, 1>(TileShape_PDV{}), _1{}));  // no mcast for O
+
+  static constexpr int VEC_SIZE = cute::ceil_div(128, sizeof_bits_v<DTypeO>); // 8
+  static_assert(HEAD_DIM_VO % VEC_SIZE == 0);
+  static constexpr int NUM_THREADS_PER_ROW = HEAD_DIM_VO / VEC_SIZE; // 64
+  static_assert(NUM_MMA_THREADS % NUM_THREADS_PER_ROW == 0);
+  static constexpr int NUM_ROWS = NUM_MMA_THREADS / NUM_THREADS_PER_ROW;
+  using TiledCopyOAtom = cute::Copy_Atom<cute::UniversalCopy<cutlass::uint128_t>, DTypeO>;
+  using TiledCopyOThrLayout = decltype(cute::make_layout(
+      cute::make_shape(Int<NUM_ROWS>{}, Int<NUM_THREADS_PER_ROW>{}), LayoutRight{}));
+  using TiledCopyOValLayout =
+      decltype(cute::make_layout(cute::make_shape(_1{}, Int<VEC_SIZE>{}), LayoutRight{}));
+  using TiledCopyO =
+      decltype(make_tiled_copy(TiledCopyOAtom{}, TiledCopyOThrLayout{},  // Thr layout
+                               TiledCopyOValLayout{}                     // Val layout
+                               ));
+  struct Arguments {
+    DTypeO* O_ptr;
+    LayoutNTMAT const layout_O;
+    DTypeO* O_ptr_tmp;
+    LayoutNTMAT const layout_O_tmp;
+  };
+
+  // Device side kernel params
+  struct Params {
+    DTypeO* O_ptr;
+    LayoutNTMAT const layout_O;
+    DTypeO* O_ptr_tmp;
+    LayoutNTMAT const layout_O_tmp;
+  };
+
+  static Params to_underlying_arguments_ntma(Arguments const& args) {
+    return {args.O_ptr, args.layout_O, args.O_ptr_tmp, args.layout_O_tmp};
+  }
+
+  CUTLASS_DEVICE
+  static void prefetch_tma_descriptors(Params const& epilogue_params) {}
+
+  template <typename SharedStorage, typename FrgTensorO, typename FrgTensorLSE,
+            typename TiledMma>
+  CUTLASS_DEVICE void store(Params const& epilogue_params, 
+                            FrgTensorO const& tOrO,
+                            FrgTensorLSE const& lse, 
+                            SharedStorage& shared_storage,
+                            TiledMma tiled_mma, 
+                            const int thread_idx,
+                            const int bid,
+                            const int bsz,
+                            const int seq_len_now,
+                            const int start_token_idx,
+                            const int tile_idx,
+                            const int q_tile_idx,
+                            const int kv_len,
+                            const int chunk_size,
+                            const int draft_total_token_num,
+                            const int o_stride_head) {
+    const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
+    const int num_chunks = cute::ceil_div(kv_len, chunk_size);
+    Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o.data()), SmemLayoutO{});
+    auto smem_tiled_copy_O = make_tiled_copy_C(SmemCopyAtomO{}, tiled_mma);
+    auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx);
+
+    Tensor tOrO_out = convert_type<DTypeO>(tOrO);
+    Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out);  // ((Atom,AtomNum), MMA_M, MMA_N)
+    Tensor taccOsO = smem_thr_copy_O.partition_D(sO);     // ((Atom,AtomNum),PIPE_M,PIPE_N)
+    // make sure gemm done
+    cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS,
+                                      /*id=*/static_cast<int>(NamedBarriers::kValueEmpty));
+    // r2s
+    cute::copy(smem_tiled_copy_O, taccOrO, taccOsO);
+    // make sure r2s done
+    cutlass::arch::NamedBarrier::sync(NUM_MMA_THREADS,
+                                      /*id=*/static_cast<int>(NamedBarriers::kValueEmpty));
+    TiledCopyO gmem_tiled_copy_O;
+    auto O_ptr = num_chunks == 1 ? epilogue_params.O_ptr + (start_token_idx * Ktraits::GROUP_SIZE + q_group_offset) * o_stride_head : epilogue_params.O_ptr_tmp + (((tile_idx * bsz + bid) * draft_total_token_num) * Ktraits::GROUP_SIZE + q_group_offset) * o_stride_head;
+    Tensor mO = make_tensor(make_gmem_ptr(O_ptr), epilogue_params.layout_O);
+    Tensor gO = local_tile(mO, select<0, 1>(TileShape_PDV{}), make_coord(_, _0{}))(_, _, _0{});
+    Tensor cO = make_identity_tensor(gO.shape());  // (O, D) -> (o_idx, d_idx)
+    ThrCopy thr_copy_O = gmem_tiled_copy_O.get_slice(thread_idx);
+    Tensor tOgO = thr_copy_O.partition_D(gO);  // (CPY, CPY_O, CPY_D)
+    Tensor tOsO = thr_copy_O.partition_S(sO);  // (CPY, CPY_O, CPY_D)
+    Tensor tOcO = thr_copy_O.partition_D(cO);  // (CPY, CPY_O, CPY_D)
+    Tensor tOgOGroup = flatten_1(tOgO);        // (CPY, (CPY_O, CPY_D))
+    Tensor tOsOGroup = flatten_1(tOsO);        // (CPY, (CPY_O, CPY_D))
+    Tensor tOcOGroup = flatten_1(tOcO);        // (CPY, (CPY_O, CPY_D))
+
+    // copy if not out of bound
+    auto predicate_fn = [&](auto coords) {
+      auto s_coords = tOcOGroup(_0{}, coords);
+      return elem_less((get<0>(s_coords) + q_group_offset) / Ktraits::GROUP_SIZE, seq_len_now);
+    };
+    copy_if(gmem_tiled_copy_O, predicate_fn, tOsOGroup, tOgOGroup);
+  }
+};
+
+}  // namespace mla_attn
+
+#endif  // ATTENTION_HOPPER_EPILOGUE_CUH_

csrc/gpu/mla_attn/kernel_traits.cuh

@@ -0,0 +1,165 @@
+// 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.
+
+/*
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
+ * Dao. Licensed under the BSD 3-Clause.
+ *
+ * Modified by the FlashInfer team.
+ */
+
+#pragma once
+
+#ifndef ATTENTION_HOPPER_KERNEL_TRAITS_CUH_
+#define ATTENTION_HOPPER_KERNEL_TRAITS_CUH_
+
+#include <type_traits>
+
+#include "cute/algorithm/copy.hpp"
+#include "cute/atom/mma_atom.hpp"
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/layout/layout.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/pipeline/pipeline.hpp"
+
+namespace mla_attn {
+
+using namespace cute;
+
+template <typename MainloopPipeline, typename MainloopPipelineQ, class DTypeQ, class DTypeKV, class DTypeQKAccum, class DTypeOut, class IdType,
+          int BLOCK_SHAPE_KV, class SmemLayoutQ, class SmemLayoutK, class SmemLayoutP, class SmemLayoutRow, class SmemLayoutO>
+struct alignas(16) SharedStorageQKVO {
+  alignas(16) cute::array_aligned<DTypeQ, cute::cosize_v<SmemLayoutQ>> smem_q;
+  alignas(16) cute::array_aligned<DTypeQ, cute::cosize_v<SmemLayoutP>> smem_p;
+  alignas(16) cute::array_aligned<DTypeQKAccum, cute::cosize_v<SmemLayoutRow>> smem_scale;
+  union {
+    alignas(16) cute::array_aligned<DTypeKV, cute::cosize_v<SmemLayoutK>> smem_kv;
+    alignas(16) cute::array_aligned<DTypeOut, cute::cosize_v<SmemLayoutO>> smem_o;
+  };
+  struct {
+    alignas(16) typename MainloopPipelineQ::SharedStorage pipeline_q;
+    alignas(16) typename MainloopPipeline::SharedStorage pipeline_kv;
+  };
+};
+
+template <bool USE_TMA_LOAD_KV_, int HEAD_DIM_QK_, int HEAD_DIM_VO_, int GROUP_SIZE_, int BLOCK_SHAPE_Q_, int BLOCK_SHAPE_KV_,
+          int NUM_STAGES_, typename DTypeQ_, typename DTypeKV_, typename DTypeO_, typename IdType_, typename NV_TYPE_>
+struct AttentionKernelTraits {
+
+  using DTypeQ = DTypeQ_;
+  using DTypeKV = DTypeKV_;
+  using DTypeO = DTypeO_;
+  using IdType = IdType_;
+  using DTypeQKAccum = float;
+  using DTypePVAccum = float;
+  using NV_TYPE = NV_TYPE_;
+  
+  
+  static constexpr bool USE_TMA_LOAD_KV = USE_TMA_LOAD_KV_;
+  static constexpr int GROUP_SIZE = GROUP_SIZE_;
+  static constexpr int BLOCK_SHAPE_Q = BLOCK_SHAPE_Q_;
+  static_assert(BLOCK_SHAPE_Q % 64 == 0);
+  static constexpr int BLOCK_SHAPE_KV = BLOCK_SHAPE_KV_;
+  static constexpr int HEAD_DIM_QK = HEAD_DIM_QK_;
+  static constexpr int HEAD_DIM_VO = HEAD_DIM_VO_;
+  static constexpr int NUM_PER_STAGE = BLOCK_SHAPE_KV * HEAD_DIM_QK;
+  static_assert(HEAD_DIM_QK % 32 == 0);
+  static_assert(HEAD_DIM_VO % 32 == 0);
+
+  static constexpr int NUM_WARPS = 12;
+  static constexpr int NUM_THREADS = 384;
+  static constexpr int NUM_PRODUCER_THREADS = 128;
+
+  using TileShape_QKD = Shape<Int<BLOCK_SHAPE_Q>, Int<BLOCK_SHAPE_KV>, Int<HEAD_DIM_QK>>;
+  using TileShape_PDV = Shape<Int<BLOCK_SHAPE_Q>, Int<HEAD_DIM_VO>, Int<BLOCK_SHAPE_KV>>;
+
+  static constexpr int NUM_STAGES = NUM_STAGES_;
+
+  using AtomLayoutQKD = Layout<Shape<Int<BLOCK_SHAPE_Q / 64>, _1, _1>>;
+  using AtomLayoutPV = Layout<Shape<Int<BLOCK_SHAPE_Q / 64>, _2, _1>>;
+  using TiledMmaQK = decltype(cute::make_tiled_mma(
+      cute::GMMA::ss_op_selector<DTypeQ, DTypeKV, DTypeQKAccum, TileShape_QKD>(), AtomLayoutQKD{}));
+  using TiledMmaPV = decltype(cute::make_tiled_mma(
+      cute::GMMA::rs_op_selector<DTypeKV, DTypeKV, /*ElementAccum=*/DTypePVAccum, TileShape_PDV,
+                                 GMMA::Major::K, GMMA::Major::MN>(),
+      AtomLayoutPV{}));
+  using TiledMmaPVSS = decltype(cute::make_tiled_mma(
+      cute::GMMA::ss_op_selector<DTypeKV, DTypeKV, /*ElementAccum=*/DTypePVAccum, TileShape_PDV,
+                                 GMMA::Major::K, GMMA::Major::MN>(),
+      AtomLayoutPV{}));
+
+  static constexpr int NUM_MMA_THREADS = size(TiledMmaPV{});
+
+  using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})),
+                                   decltype(cute::get<2>(TileShape_QKD{}))>());
+  using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_QKD{})));
+
+  using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeKV, decltype(cute::get<1>(TileShape_QKD{})),
+                                   decltype(cute::get<2>(TileShape_QKD{}))>());
+  using SmemLayoutK = decltype(tile_to_shape(
+      SmemLayoutAtomK{},
+      make_shape(shape<1>(TileShape_QKD{}), shape<2>(TileShape_QKD{}), Int<NUM_STAGES>{})));
+  using SmemLayoutVt = decltype(composition(
+      SmemLayoutK{}, make_ordered_layout(make_shape(get<2>(TileShape_QKD{}),
+                                                    get<1>(TileShape_QKD{}), Int<NUM_STAGES>{}),
+                                         Step<_2, _1, _3>{})));
+  using SmemLayoutAtomV = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeKV, decltype(cute::get<2>(TileShape_PDV{})),
+                                   decltype(cute::get<1>(TileShape_PDV{}))>());
+  using SmemLayoutV = decltype(tile_to_shape(
+      SmemLayoutAtomV{},
+      make_shape(get<2>(TileShape_PDV{}), get<1>(TileShape_PDV{}), Int<1>{})));
+
+  // Note this is the transpose in terms of the view, not in terms of memory.
+  using SmemLayoutVtOneStage = decltype(composition(
+      SmemLayoutV{}, make_ordered_layout(make_shape(get<1>(TileShape_PDV{}),
+                                                    get<2>(TileShape_PDV{}), Int<1>{}),
+                                         Step<_2, _1, _3>{})));
+
+  using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_PDV{})),
+                                   decltype(cute::get<1>(TileShape_PDV{}))>());
+  using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{})));
+
+  using SmemCopyAtom = Copy_Atom<cute::SM90_U32x4_STSM_N, DTypeQ>;
+  
+  static constexpr bool IS_CTA_32 = (BLOCK_SHAPE_KV == 32);
+  using SmemLayoutRowOneStage = Layout<Shape<_2, Int<128>>, Stride<_1, _2>>;
+  using SmemLayoutRowTwoStage = Layout<Shape<_2, Int<128>, _2>, Stride<_1, _2, _256>>;
+  using SmemLayoutRow = std::conditional_t<IS_CTA_32, SmemLayoutRowTwoStage, SmemLayoutRowOneStage>;
+
+  using SmemLayoutAtomP = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})),
+                                   decltype(cute::get<1>(TileShape_QKD{}))>());
+  using SmemLayoutPSSOneStage = decltype(tile_to_shape(SmemLayoutAtomP{}, select<0, 1>(TileShape_QKD{})));
+  using SmemLayoutPSSTwoStage = decltype(tile_to_shape(SmemLayoutAtomP{}, make_shape(Int<BLOCK_SHAPE_Q>{}, Int<BLOCK_SHAPE_KV>{}, Int<2>{})));
+  using SmemLayoutP = std::conditional_t<IS_CTA_32, SmemLayoutPSSTwoStage, SmemLayoutPSSOneStage>;
+
+  using MainloopPipelineQ = typename cutlass::PipelineAsync<1>;
+  using PipelineStateQ = typename cutlass::PipelineState<1>;
+  using MainloopPipeline =
+      std::conditional_t<USE_TMA_LOAD_KV, typename cutlass::PipelineTmaAsync<NUM_STAGES>,
+                         typename cutlass::PipelineAsync<NUM_STAGES>>;
+  using PipelineState = typename cutlass::PipelineState<NUM_STAGES>;
+
+  using SharedStorage = SharedStorageQKVO<MainloopPipeline, MainloopPipelineQ, DTypeQ, DTypeKV, DTypeQKAccum, DTypeO, IdType, BLOCK_SHAPE_KV,
+                                          SmemLayoutQ, SmemLayoutK, SmemLayoutP, SmemLayoutRow, SmemLayoutO>;
+};
+
+}  // namespace mla_attn
+
+#endif

csrc/gpu/mla_attn/kernel_traits_wg4.cuh

@@ -0,0 +1,160 @@
+// 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
+
+#ifndef ATTENTION_HOPPER_KERNEL_TRAITS_WG4_CUH_
+#define ATTENTION_HOPPER_KERNEL_TRAITS_WG4_CUH_
+
+#include <type_traits>
+#include "cute/algorithm/copy.hpp"
+#include "cute/atom/mma_atom.hpp"
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/layout/layout.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/pipeline/pipeline.hpp"
+
+namespace mla_attn {
+
+using namespace cute;
+
+template <typename MainloopPipeline, typename MainloopPipelineQ, typename MainloopPipelineQK, class DTypeQ, class DTypeKV, class DTypeQKAccum, class DTypeOut, class IdType,
+          int BLOCK_SHAPE_KV, class SmemLayoutQ, class SmemLayoutK, class SmemLayoutP, class SmemLayoutRow, class SmemLayoutO>
+struct alignas(16) SharedStorageQKVWG4O {
+  alignas(16) cute::array_aligned<DTypeQ, cute::cosize_v<SmemLayoutQ>> smem_q;
+  alignas(16) cute::array_aligned<DTypeQ, cute::cosize_v<SmemLayoutP>> smem_p;
+  alignas(16) cute::array_aligned<DTypeQKAccum, cute::cosize_v<SmemLayoutRow>> smem_scale;
+  union {
+    alignas(16) cute::array_aligned<DTypeKV, cute::cosize_v<SmemLayoutK>> smem_kv;
+    alignas(16) cute::array_aligned<DTypeOut, cute::cosize_v<SmemLayoutO>> smem_o;
+  };
+  struct {
+    alignas(16) typename MainloopPipelineQ::SharedStorage pipeline_q;
+    alignas(16) typename MainloopPipelineQK::SharedStorage pipeline_qk;
+    alignas(16) typename MainloopPipeline::SharedStorage pipeline_kv;
+  };
+};
+
+template <bool USE_TMA_LOAD_KV_, int HEAD_DIM_QK_, int HEAD_DIM_VO_, int GROUP_SIZE_, int BLOCK_SHAPE_Q_, int BLOCK_SHAPE_KV_,
+          int NUM_STAGES_, typename DTypeQ_, typename DTypeKV_, typename DTypeO_, typename IdType_, typename NV_TYPE_>
+struct AttentionKernelWG4Traits {
+
+  using DTypeQ = DTypeQ_;
+  using DTypeKV = DTypeKV_;
+  using DTypeO = DTypeO_;
+  using IdType = IdType_;
+//   using DTypeQKAccum = float;
+//   using DTypePVAccum = float;
+  using DTypeQKAccum = DTypeQ;
+  using DTypePVAccum = DTypeO;
+  using NV_TYPE = NV_TYPE_;
+  
+  
+  static constexpr bool USE_TMA_LOAD_KV = USE_TMA_LOAD_KV_;
+  static constexpr int GROUP_SIZE = GROUP_SIZE_;
+  static constexpr int BLOCK_SHAPE_Q = BLOCK_SHAPE_Q_;
+  static_assert(BLOCK_SHAPE_Q % 64 == 0);
+  static constexpr int BLOCK_SHAPE_KV = BLOCK_SHAPE_KV_;
+  static constexpr int HEAD_DIM_QK = HEAD_DIM_QK_;
+  static constexpr int HEAD_DIM_VO = HEAD_DIM_VO_;
+  static constexpr int NUM_PER_STAGE = BLOCK_SHAPE_KV * HEAD_DIM_QK;
+  static_assert(HEAD_DIM_QK % 32 == 0);
+  static_assert(HEAD_DIM_VO % 32 == 0);
+
+
+  using TileShape_QKD = Shape<Int<BLOCK_SHAPE_Q>, Int<BLOCK_SHAPE_KV>, Int<HEAD_DIM_QK>>;
+  using TileShape_PDV = Shape<Int<BLOCK_SHAPE_Q>, Int<HEAD_DIM_VO>, Int<BLOCK_SHAPE_KV>>;
+
+  static constexpr int NUM_STAGES = NUM_STAGES_;
+
+  using AtomLayoutQKD = Layout<Shape<Int<BLOCK_SHAPE_Q / 64>, _1, _1>>;
+  using AtomLayoutPV = Layout<Shape<Int<BLOCK_SHAPE_Q / 64>, _2, _1>>;
+  using TiledMmaQK = decltype(cute::make_tiled_mma(
+      cute::GMMA::ss_op_selector<DTypeQ, DTypeKV, DTypeQKAccum, TileShape_QKD>(), AtomLayoutQKD{}));
+  using TiledMmaPV = decltype(cute::make_tiled_mma(
+      cute::GMMA::rs_op_selector<DTypeKV, DTypeKV, /*ElementAccum=*/DTypePVAccum, TileShape_PDV,
+                                 GMMA::Major::K, GMMA::Major::MN>(),
+      AtomLayoutPV{}));
+  using TiledMmaPVSS = decltype(cute::make_tiled_mma(
+      cute::GMMA::ss_op_selector<DTypeKV, DTypeKV, /*ElementAccum=*/DTypePVAccum, TileShape_PDV,
+                                 GMMA::Major::K, GMMA::Major::MN>(),
+      AtomLayoutPV{}));
+
+  static constexpr int NUM_MMA_THREADS_QK = size(TiledMmaQK{});
+  static constexpr int NUM_MMA_THREADS_PV = size(TiledMmaPVSS{});
+  static constexpr int NUM_MMA_THREADS = NUM_MMA_THREADS_PV;
+  static constexpr int NUM_PRODUCER_THREADS = cutlass::NumThreadsPerWarpGroup;
+  static constexpr int NUM_THREADS = NUM_MMA_THREADS_QK + NUM_MMA_THREADS_PV + NUM_PRODUCER_THREADS;
+  static constexpr int NUM_WARPS = NUM_THREADS / cutlass::NumThreadsPerWarp;
+
+  using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})),
+                                   decltype(cute::get<2>(TileShape_QKD{}))>());
+  using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShape_QKD{})));
+
+  using SmemLayoutAtomK = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeKV, decltype(cute::get<1>(TileShape_QKD{})),
+                                   decltype(cute::get<2>(TileShape_QKD{}))>());
+  using SmemLayoutK = decltype(tile_to_shape(
+      SmemLayoutAtomK{},
+      make_shape(shape<1>(TileShape_QKD{}), shape<2>(TileShape_QKD{}), Int<NUM_STAGES>{})));
+  using SmemLayoutVt = decltype(composition(
+      SmemLayoutK{}, make_ordered_layout(make_shape(get<2>(TileShape_QKD{}),
+                                                    get<1>(TileShape_QKD{}), Int<NUM_STAGES>{}),
+                                         Step<_2, _1, _3>{})));
+  using SmemLayoutAtomV = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeKV, decltype(cute::get<2>(TileShape_PDV{})),
+                                   decltype(cute::get<1>(TileShape_PDV{}))>());
+  using SmemLayoutV = decltype(tile_to_shape(
+      SmemLayoutAtomV{},
+      make_shape(get<2>(TileShape_PDV{}), get<1>(TileShape_PDV{}), Int<1>{})));
+
+  // Note this is the transpose in terms of the view, not in terms of memory.
+  using SmemLayoutVtOneStage = decltype(composition(
+      SmemLayoutV{}, make_ordered_layout(make_shape(get<1>(TileShape_PDV{}),
+                                                    get<2>(TileShape_PDV{}), Int<1>{}),
+                                         Step<_2, _1, _3>{})));
+
+  using SmemLayoutAtomO = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeO, decltype(cute::get<0>(TileShape_PDV{})),
+                                   decltype(cute::get<1>(TileShape_PDV{}))>());
+  using SmemLayoutO = decltype(tile_to_shape(SmemLayoutAtomO{}, select<0, 1>(TileShape_PDV{})));
+
+  using SmemCopyAtom = Copy_Atom<cute::SM90_U32x4_STSM_N, DTypeQ>;
+  // todo: permute layout
+  using SmemLayoutRow = Layout<Shape<Int<2>, Int<128>, Int<2>>, Stride<Int<1>, Int<2>, Int<256>>>;
+
+  using SmemLayoutAtomP = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
+                                   GMMA::Major::K, DTypeQ, decltype(cute::get<0>(TileShape_QKD{})),
+                                   decltype(cute::get<1>(TileShape_QKD{}))>());
+  static constexpr int QK_STAGE = 2;
+  using SmemLayoutPSS = decltype(tile_to_shape(SmemLayoutAtomP{}, make_shape(shape<0>(TileShape_QKD{}), shape<1>(TileShape_QKD{}), Int<QK_STAGE>{})));
+
+  using MainloopPipelineQ = typename cutlass::PipelineAsync<1>;
+  using PipelineStateQ = typename cutlass::PipelineState<1>;
+  using MainloopPipelineQK = typename cutlass::PipelineAsync<QK_STAGE>;
+  using PipelineStateQK = typename cutlass::PipelineState<QK_STAGE>;
+  using MainloopPipeline =
+      std::conditional_t<USE_TMA_LOAD_KV, typename cutlass::PipelineTmaAsync<NUM_STAGES>,
+                         typename cutlass::PipelineAsync<NUM_STAGES>>;
+  using PipelineState = typename cutlass::PipelineState<NUM_STAGES>;
+
+  using SharedStorage = SharedStorageQKVWG4O<MainloopPipeline, MainloopPipelineQ, MainloopPipelineQK, DTypeQ, DTypeKV, DTypeQKAccum, DTypeO, IdType, BLOCK_SHAPE_KV,
+                                             SmemLayoutQ, SmemLayoutK, SmemLayoutPSS, SmemLayoutRow, SmemLayoutO>;
+};
+
+}  // namespace mla_attn
+
+#endif // ATTENTION_HOPPER_KERNEL_TRAITS_WG4_CUH_

csrc/gpu/mla_attn/mainloop_load.cuh

@@ -0,0 +1,355 @@
+// 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
+
+#ifndef ATTENTION_HOPPER_MAINLOOP_LOAD_CUH_
+#define ATTENTION_HOPPER_MAINLOOP_LOAD_CUH_
+
+#include <cutlass/array.h>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/numeric_types.h>
+
+#include "cute/tensor.hpp"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/pipeline/pipeline.hpp"
+#include "named_barrier.cuh"
+#include "utils.cuh"
+
+#ifdef DEBUG_MLA
+#undef DEBUG_MLA
+#endif
+// #define DEBUG_MLA
+
+namespace mla_attn {
+
+using namespace cute;
+
+template <typename Ktraits, bool CAUSAL>
+struct CollectiveMainloop {
+  using DTypeQ = typename Ktraits::DTypeQ;
+  using DTypeKV = typename Ktraits::DTypeKV;
+  using DTypeMD = float;
+  using IdType = typename Ktraits::IdType;
+  using TileShape_QKD = typename Ktraits::TileShape_QKD;
+  using TileShape_PDV = typename Ktraits::TileShape_PDV;
+  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
+  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
+
+  static constexpr int NUM_STAGES = Ktraits::NUM_STAGES;
+  static constexpr int HEAD_DIM_QK = Ktraits::HEAD_DIM_QK;
+  static constexpr int HEAD_DIM_VO = Ktraits::HEAD_DIM_VO;
+
+  using GmemTiledCopyKV = cute::SM90_TMA_LOAD;
+  static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(DTypeQ); // 8
+  static_assert(HEAD_DIM_QK % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad"); // 576 512
+  static constexpr int kGmemThreadsPerRow = 64 / kGmemElemsPerLoad; // 8
+  using AlignmentTypeQ = cute::uint_byte_t<static_cast<int>(sizeof(DTypeQ)) * kGmemElemsPerLoad>;
+  using GmemCopyAtomQ = cute::Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<AlignmentTypeQ>, DTypeQ>;
+  static constexpr int kNThreadsLoad = Ktraits::NUM_PRODUCER_THREADS;
+  static_assert(kNThreadsLoad % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow");
+  using GmemLayoutAtom = Layout<
+            Shape<Int<kNThreadsLoad / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>, // 32, 8
+            Stride<Int<kGmemThreadsPerRow>, _1>>;
+  using GmemTiledCopy = decltype(make_tiled_copy(
+          GmemCopyAtomQ{},
+          GmemLayoutAtom{},
+          Layout<Shape<_1, _8>>{}));  // Val layout, 8 vals per read
+
+  using GmemLayoutAtomQ = Layout<
+            Shape<Int<Ktraits::NUM_PRODUCER_THREADS / kGmemThreadsPerRow>, Int<kGmemThreadsPerRow>>, // 32, 8
+            Stride<Int<kGmemThreadsPerRow>, _1>>;
+  using GmemTiledCopyQ = decltype(make_tiled_copy(
+          GmemCopyAtomQ{},
+          GmemLayoutAtomQ{},
+          Layout<Shape<_1, _8>>{}));  // Val layout, 8 vals per read
+
+  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
+  using SmemLayoutAtomQ = typename Ktraits::SmemLayoutAtomQ;
+
+  using SmemLayoutK = typename Ktraits::SmemLayoutK;
+  using SmemLayoutV = typename Ktraits::SmemLayoutV;
+  using SmemLayoutVt = typename Ktraits::SmemLayoutVt;
+
+  using ShapeQT = cute::Shape<int32_t, int32_t>;
+  using StrideQT = cute::Shape<int32_t, _1>;
+  using LayoutQT = cute::Layout<ShapeQT, StrideQT>;
+
+  using ShapeT = cute::Shape<int32_t, int32_t, int32_t>;
+  using StrideT = cute::Shape<int32_t, _1, int32_t>;
+  using LayoutT = cute::Layout<ShapeT, StrideT>;
+
+  using ShapeMDT = cute::Shape<int32_t, int32_t>;
+  using StrideMDT = cute::Shape<int32_t, _1>;
+  using LayoutMDT = cute::Layout<ShapeMDT, StrideMDT>;
+  
+  using TMA_KV = decltype(make_tma_copy(
+      GmemTiledCopyKV{},
+      make_tensor(
+          make_gmem_ptr(static_cast<DTypeKV const*>(nullptr)), 
+          repeat_like(StrideT{}, int32_t(0)), StrideT{}
+      ),
+      take<0, 2>(SmemLayoutK{}),
+      select<1, 2>(TileShape_QKD{}),
+      _1{})); // no mcast for KV
+
+  static constexpr bool USE_TMA_LOAD_KV = Ktraits::USE_TMA_LOAD_KV;
+  using MainloopPipeline = typename Ktraits::MainloopPipeline;
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState = typename MainloopPipeline::PipelineState;
+
+  using MainloopPipelineQ = typename Ktraits::MainloopPipelineQ;
+  using PipelineParamsQ = typename MainloopPipelineQ::Params;
+  using PipelineStateQ = typename MainloopPipelineQ::PipelineState;
+
+  static constexpr uint32_t TmaTransactionBytesQ =
+      static_cast<uint32_t>(size(SmemLayoutQ{}) * cutlass::sizeof_bits_v<DTypeQ> / 8);
+  static constexpr uint32_t TmaTransactionBytesKV =
+      static_cast<uint32_t>(size(take<0, 2>(SmemLayoutK{})) * cutlass::sizeof_bits_v<DTypeKV> / 8);
+
+  // Host side kernel arguments
+  struct Arguments {
+    LayoutQT layout_Q;
+    LayoutT layout_KV;
+    LayoutMDT layout_MD;
+    DTypeQ const* Q_ptr;
+    DTypeKV const* KV_ptr;
+    DTypeMD const* m_ptr;
+    DTypeMD const* d_ptr;
+    IdType const* kv_block_tables;
+    IdType const* seq_lens_this_time;
+    IdType const* seq_lens_encoder;
+    IdType const* seq_lens_decoder;
+    IdType const* cumsum_q_seqlens;
+    IdType const* batch_ids;
+    IdType const* tile_ids_per_batch;
+    IdType const* q_tile_ids_per_batch;
+    IdType const* num_blocks_x;
+    float sm_scale;
+    int bsz;
+    int max_block_num;
+    int max_block_num_per_seq;
+    int q_stride_bsz;
+    int q_stride_head_num;
+    int kv_stride_block_num;
+    int kv_stride_block_size;
+    int o_stride_bsz;
+    int o_stride_head_num;
+    int chunk_size;
+    int chunk_num;
+    int draft_total_token_num;
+  };
+
+  // Device side kernel params
+  struct Params {
+    LayoutQT layout_Q;
+    LayoutT layout_KV;
+    LayoutMDT layout_MD;
+    DTypeQ *Q_ptr;
+    DTypeKV* KV_ptr;
+    DTypeMD* m_ptr;
+    DTypeMD* d_ptr;
+    IdType* kv_block_tables;
+    IdType* seq_lens_this_time;
+    IdType* seq_lens_encoder;
+    IdType* seq_lens_decoder;
+    IdType* cumsum_q_seqlens;
+    IdType* batch_ids;
+    IdType* tile_ids_per_batch;
+    IdType* q_tile_ids_per_batch;
+    IdType* num_blocks_x;
+    float sm_scale;
+    int bsz;
+    int max_block_num;
+    int max_block_num_per_seq;
+    int q_stride_bsz;
+    int q_stride_head_num;
+    int kv_stride_block_num;
+    int kv_stride_block_size;
+    int o_stride_bsz;
+    int o_stride_head_num;
+    int chunk_size;
+    int chunk_num;
+    int draft_total_token_num;
+    TMA_KV tma_load_KV;
+  };
+
+  static Params to_underlying_arguments(Arguments const& args) {
+    TMA_KV tma_load_KV;
+    if constexpr (USE_TMA_LOAD_KV) {
+      Tensor mKV = make_tensor(make_gmem_ptr(args.KV_ptr), args.layout_KV);
+      tma_load_KV = 
+          make_tma_copy(GmemTiledCopyKV{}, mKV, SmemLayoutK{}(_, _, _0{}), select<1, 2>(TileShape_QKD{}), _1{});
+    }
+    return {args.layout_Q,
+            args.layout_KV,
+            args.layout_MD,
+            const_cast<DTypeQ*>(args.Q_ptr),
+            const_cast<DTypeKV*>(args.KV_ptr),
+            const_cast<DTypeMD*>(args.m_ptr),
+            const_cast<DTypeMD*>(args.d_ptr),
+            const_cast<IdType*>(args.kv_block_tables),
+            const_cast<IdType*>(args.seq_lens_this_time),
+            const_cast<IdType*>(args.seq_lens_encoder),
+            const_cast<IdType*>(args.seq_lens_decoder),
+            const_cast<IdType*>(args.cumsum_q_seqlens),
+            const_cast<IdType*>(args.batch_ids),
+            const_cast<IdType*>(args.tile_ids_per_batch),
+            const_cast<IdType*>(args.q_tile_ids_per_batch),
+            const_cast<IdType*>(args.num_blocks_x),
+            args.sm_scale,
+            args.bsz,
+            args.max_block_num,
+            args.max_block_num_per_seq,
+            args.q_stride_bsz,
+            args.q_stride_head_num,
+            args.kv_stride_block_num,
+            args.kv_stride_block_size,
+            args.o_stride_bsz,
+            args.o_stride_head_num,
+            args.chunk_size,
+            args.chunk_num,
+            args.draft_total_token_num,
+            tma_load_KV
+            };
+  }
+
+  CUTLASS_DEVICE
+  static void prefetch_tma_descriptors(Params const& mainloop_params) {
+    if constexpr (USE_TMA_LOAD_KV) {
+      cute::prefetch_tma_descriptor(mainloop_params.tma_load_KV.get_tma_descriptor());
+    }
+  }
+
+  template <typename SharedStorage>
+  CUTLASS_DEVICE void load_q(Params const& mainloop_params, 
+                             MainloopPipelineQ pipeline_q,
+                             PipelineStateQ& smem_pipe_write_q,
+                             SharedStorage& shared_storage,
+                             const int thread_idx,
+                             const int bid,
+                             const int q_tile_idx) {
+    const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
+    int start_q_head_idx = mainloop_params.cumsum_q_seqlens[bid] * Ktraits::GROUP_SIZE + q_group_offset;
+    int offset_Q = mainloop_params.q_stride_head_num * start_q_head_idx;
+    Tensor mQ = make_tensor(make_gmem_ptr(mainloop_params.Q_ptr + offset_Q), mainloop_params.layout_Q);
+    Tensor gQ =
+        local_tile(mQ, select<0, 2>(TileShape_QKD{}), make_coord(_, _0{}))(_, _, _0{});
+    Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+    Tensor cQ = cute::make_identity_tensor(gQ.shape());
+
+    GmemTiledCopyQ gmem_tiled_copy_q;
+    auto gmem_thr_copy_q = gmem_tiled_copy_q.get_slice(thread_idx);
+    Tensor tQgQ = gmem_thr_copy_q.partition_S(gQ);
+    Tensor tQsQ = gmem_thr_copy_q.partition_D(sQ);
+    Tensor tQcQ = gmem_thr_copy_q.partition_D(cQ);
+    Tensor tQcQGroup = flatten_1(tQcQ);
+
+    int valid_q_size = mainloop_params.seq_lens_this_time[bid];
+    auto q_predicate_fn = [&](auto coords) {
+      auto s_coords = tQcQGroup(_0{}, coords);
+      return elem_less((get<0>(s_coords) + q_group_offset) / Ktraits::GROUP_SIZE, valid_q_size);
+    };
+    Tensor tQgQiGroup = flatten_1(tQgQ);
+    Tensor tQsQiGroup = flatten_1(tQsQ);
+
+    pipeline_q.producer_acquire(smem_pipe_write_q);
+    copy_if(gmem_tiled_copy_q, q_predicate_fn, tQgQiGroup, tQsQiGroup);
+    pipeline_q.producer_commit(smem_pipe_write_q, cutlass::arch::cpasync_barrier_arrive);
+    ++smem_pipe_write_q;
+  }
+
+  template <typename SharedStorage>
+  CUTLASS_DEVICE void load_kv(Params const& mainloop_params, 
+                              MainloopPipeline pipeline_kv,
+                              PipelineState& smem_pipe_write_kv,
+                              SharedStorage& shared_storage,
+                              const int bid,
+                              const int kv_len,
+                              const int tile_idx) {
+    int thread_idx = threadIdx.x;
+    int warp_idx_in_warpgroup = __shfl_sync(0xffffffff, (thread_idx / 32) % 4, 0);
+    Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+    Tensor mKV = make_tensor(make_gmem_ptr(mainloop_params.KV_ptr), mainloop_params.layout_KV);
+    Tensor gKV = local_tile(mKV, make_shape(get<1>(TileShape_QKD{}), get<2>(TileShape_QKD{})), make_coord(_, _))(_, _, _0{}, _0{}, _);
+    GmemTiledCopy gmem_tiled_copy_kv;
+    auto gmem_thr_copy_kv = gmem_tiled_copy_kv.get_slice(thread_idx);
+
+    static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
+    const int start_len = tile_idx * mainloop_params.chunk_size;
+    const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
+    const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
+
+    auto kv_block_tables = make_tensor(make_gmem_ptr(mainloop_params.kv_block_tables), make_layout(make_shape(mainloop_params.bsz, mainloop_params.max_block_num_per_seq), make_stride(mainloop_params.max_block_num_per_seq, 1)));
+
+    Tensor tKgK = gmem_thr_copy_kv.partition_S(gKV);
+    Tensor tKsK = gmem_thr_copy_kv.partition_S(sK);
+
+    for (int kv_tile_idx = end_tile_idx; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
+      const int block_idx = kv_block_tables(bid, kv_tile_idx);
+      pipeline_kv.producer_acquire(smem_pipe_write_kv);
+      Tensor tKgKiGroup = flatten_1(tKgK(_, _, _, block_idx));
+      Tensor tKsKiGroup =
+          flatten_1(tKsK(_, _, _, smem_pipe_write_kv.index()));
+      copy(gmem_tiled_copy_kv, tKgKiGroup, tKsKiGroup);
+      pipeline_kv.producer_commit(smem_pipe_write_kv, cutlass::arch::cpasync_barrier_arrive);
+      ++smem_pipe_write_kv;
+    }
+  }
+
+  template <typename SharedStorage>
+  CUTLASS_DEVICE void load_kv_tma(Params const& mainloop_params, 
+                                  MainloopPipeline pipeline_kv,
+                                  PipelineState& smem_pipe_write_kv,
+                                  SharedStorage& shared_storage,
+                                  const int bid,
+                                  const int kv_len,
+                                  const int tile_idx) {
+    int thread_idx = threadIdx.x;
+    Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+
+    Tensor mKV = mainloop_params.tma_load_KV.get_tma_tensor(mainloop_params.layout_KV.shape());
+
+    // Prepare the TMA loads
+    Tensor gKV = local_tile(mKV, make_shape(get<1>(TileShape_QKD{}), get<2>(TileShape_QKD{})), make_coord(_, _))(_, _, _0{}, _0{}, _);
+    auto [tKgK, tKsK] = 
+        tma_partition(mainloop_params.tma_load_KV, _0{}, Layout<_1>{},
+                      group_modes<0, 2>(sK), group_modes<0, 2>(gKV));
+
+    static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
+    const int start_len = tile_idx * mainloop_params.chunk_size;
+    const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
+    const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
+
+    auto kv_block_tables = make_tensor(make_gmem_ptr(mainloop_params.kv_block_tables), make_layout(make_shape(mainloop_params.bsz, mainloop_params.max_block_num_per_seq), make_stride(mainloop_params.max_block_num_per_seq, 1)));
+
+    int lane_predicate = cute::elect_one_sync();
+
+    if (lane_predicate) {
+#pragma unroll 2
+      for (int kv_tile_idx = end_tile_idx; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
+        const int block_idx = kv_block_tables(bid, kv_tile_idx);
+        pipeline_kv.producer_acquire(smem_pipe_write_kv);
+        copy(mainloop_params.tma_load_KV.with(*pipeline_kv.producer_get_barrier(smem_pipe_write_kv), /*mcast_mask=*/0),
+             tKgK(_, block_idx), tKsK(_, smem_pipe_write_kv.index()));
+        ++smem_pipe_write_kv;
+      }
+    }
+  }
+};
+
+}  // namespace mla_attn
+
+#endif  // ATTENTION_HOPPER_SPARSE_MAINLOOP_CUH_

csrc/gpu/mla_attn/mainloop_mma.cuh

@@ -0,0 +1,962 @@
+// 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
+
+#ifndef ATTENTION_HOPPER_MAINLOOP_MMA_CUH_
+#define ATTENTION_HOPPER_MAINLOOP_MMA_CUH_
+
+#include <cutlass/array.h>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/numeric_types.h>
+#include "named_barrier.cuh"
+
+// #define DEBUG_MLA
+
+namespace mla_attn {
+
+template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQ,
+          typename PipelineState, typename PipelineStateQ, typename SharedStorage, typename FrgTensorO, typename AttentionUpdater>
+CUTLASS_DEVICE void mma_f16(const Params& mainloop_params,
+                            MainloopPipelineQ pipeline_q,
+                            PipelineStateQ& smem_pipe_read_q,
+                            MainloopPipeline pipeline_kv,
+                            PipelineState& smem_pipe_read_kv,
+                            FrgTensorO& tOrO, 
+                            AttentionUpdater& attention_updater,
+                            const int thread_idx, 
+                            const int bid,
+                            const int kv_len,
+                            const int qo_len,
+                            const int tile_idx,
+                            const int q_tile_idx,
+                            SharedStorage& shared_storage) {
+  using DTypeQ = typename Ktraits::DTypeQ;
+  using DTypeKV = typename Ktraits::DTypeKV;
+  using DTypeMD = typename Ktraits::DTypeO;
+  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
+  using IdType = typename Ktraits::IdType;
+  using TileShape_QKD = typename Ktraits::TileShape_QKD;
+  static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS;
+  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
+  using SmemLayoutK = typename Ktraits::SmemLayoutK;
+  using SmemLayoutV = typename Ktraits::SmemLayoutV;
+  using SmemLayoutP = typename Ktraits::SmemLayoutP;
+  using SmemLayoutRow = typename Ktraits::SmemLayoutRow;
+  using SmemCopyAtom = typename Ktraits::SmemCopyAtom;
+  using SmemLayoutVt = typename Ktraits::SmemLayoutVt;
+  using SmemLayoutVtOneStage = typename Ktraits::SmemLayoutVtOneStage;
+  static_assert(is_rmem<FrgTensorO>::value, "O tensor must be rmem resident.");
+
+  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size);
+
+  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
+  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
+
+  const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
+
+  Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+  Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+  // todo: split rope/norope for sVt, not use s1~s2.
+  Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutVtOneStage{});
+  Tensor sVt_s2 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
+  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutP{});
+  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
+  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr), mainloop_params.layout_MD)(tile_idx, _); // (bsz * draft_token_num * num_head)
+  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr), mainloop_params.layout_MD)(tile_idx, _);
+
+  typename Ktraits::TiledMmaQK tiled_mma_qk;
+  auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx);
+  auto smem_tiled_copy_P = make_tiled_copy_C(SmemCopyAtom{}, tiled_mma_qk);
+  auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(thread_idx);
+  Tensor tPsP = smem_thr_copy_P.partition_D(sPSS);
+  Tensor tScalesScale = s_scale(_, thread_idx % cutlass::NumThreadsPerWarpGroup);
+
+  typename Ktraits::TiledMmaPVSS tiled_mma_pv_ss;
+  auto threadMmaPVSS = tiled_mma_pv_ss.get_thread_slice(thread_idx);
+  Tensor tOrV1 = threadMmaPVSS.partition_fragment_B(sVt_s1);
+  Tensor tOrV2 = threadMmaPVSS.partition_fragment_B(sVt_s2);
+  Tensor tOrP_CS2 = threadMmaPVSS.partition_fragment_A(sPSS);
+
+  const int start_len = tile_idx * mainloop_params.chunk_size;
+  const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
+  const int end_tile_idx =cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
+  int kv_tile_idx = end_tile_idx;
+
+  auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
+    auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+    pipeline.consumer_wait(smem_pipe_read, barrier_token);
+  };
+
+  int warp_group_idx = cutlass::canonical_warp_group_idx();
+  if (warp_group_idx == 1) {
+    // consumer 0, compute qk
+    Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ);
+    Tensor tSrK = threadMmaQK.partition_fragment_B(sK);
+
+    constexpr int n_masking_steps = !CAUSAL ? 1 : cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) + 1;
+    auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; };
+    bool is_first_step = true;
+    // wait q
+    consumer_wait(pipeline_q, smem_pipe_read_q);
+    Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+#pragma unroll 1
+    for (int masking_step = n_masking_steps; kv_tile_idx >= start_tile_idx; --masking_step, --kv_tile_idx) {
+      // wait kv
+      consumer_wait(pipeline_kv, smem_pipe_read_kv);
+      // gemm qk
+      gemm</*init=*/true, /*wg_wait=*/0>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
+                                         tSrS);
+      // mask
+      if (masking_step > 0) {
+        Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
+        Tensor tScS = threadMmaQK.partition_C(cS);
+#pragma unroll
+        for (int i = 0; i < size(tSrS); ++i) {
+          int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
+          int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
+          if constexpr (!CAUSAL) {  // Just masking based on col
+            if (kv_idx >= kv_len) {
+              tSrS(i) = AttentionUpdater::fill_value;
+            }
+          } else {
+            if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) {
+              tSrS(i) = AttentionUpdater::fill_value;
+            }
+          }
+        }
+      }
+
+      // update s (exp(s - m))
+      Tensor scale_o = is_first_step ? attention_updater.update</*init=*/true>(tSrS) : attention_updater.update</*init=*/false>(tSrS);
+      is_first_step = false;
+
+      Tensor convert_tSrS = convert_type<DTypeKV>(tSrS);
+      Tensor tPrP = smem_thr_copy_P.retile_S(convert_tSrS);
+
+      // gather qk gemm res
+      cute::copy(smem_tiled_copy_P, tPrP, tPsP);
+      cute::copy(scale_o, tScalesScale);
+      // r2s fence wgmma
+      cutlass::arch::fence_view_async_shared();
+      // make sure r2s all done
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+
+      attention_updater.rescale_o(tOrO, scale_o);
+
+      // pv gemm
+      if (smem_pipe_read_kv.index() == 0) {
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2,
+                                            tOrV1(_, _, _, _0{}), tOrO);
+      } else {
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2,
+                                            tOrV2(_, _, _, _0{}), tOrO);
+      }
+
+      pipeline_kv.consumer_release(smem_pipe_read_kv);
+      ++smem_pipe_read_kv;
+      // sync WG1 WG2
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2Sync));
+    } 
+    // release q
+    pipeline_q.consumer_release(smem_pipe_read_q);
+    ++smem_pipe_read_q;
+
+    // normalize
+    Tensor scale_o = attention_updater.finalize(tSrS); // warp reduce row sum
+    if (chunk_num_this_seq == 1) {
+      // norm
+      cute::copy(scale_o, tScalesScale);
+
+      cutlass::arch::NamedBarrier::arrive(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG2));
+      attention_updater.rescale_o(tOrO, scale_o);
+    }
+
+    // WG1 write m,d back to gmem
+    if (chunk_num_this_seq > 1 && thread_idx % 4 == 0) { // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
+      const int warp_idx = thread_idx / 32;
+#pragma unroll
+      for (int w_i = 0; w_i < 2; ++w_i) {
+        const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i + q_group_offset;
+        const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
+
+        if (token_idx < qo_len) {
+          // const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
+          const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE;
+          const int write_idx = bid * bid_offset + token_group_idx;
+          mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
+          mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
+        }
+      }
+    }
+  } else if (warp_group_idx == 2) {
+    // consumer 1, compute pv
+    Tensor scale_o = make_tensor<DTypeQKAccum>(Shape<_2>{});
+    for (; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
+      // wait kv
+      consumer_wait(pipeline_kv, smem_pipe_read_kv);
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+
+      // A: tPsP
+      cute::copy(tScalesScale, scale_o);
+
+      // rescale
+      attention_updater.rescale_o(tOrO, scale_o);
+      if (smem_pipe_read_kv.index() == 0) {
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2,
+                                            tOrV1(_, _, _, _0{}), tOrO);
+      } else {
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2,
+                                            tOrV2(_, _, _, _0{}), tOrO);
+      }
+
+      pipeline_kv.consumer_release(smem_pipe_read_kv);
+      ++smem_pipe_read_kv;
+      // sync WG1 WG2
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2Sync));
+    }
+    if (chunk_num_this_seq == 1) {
+      // norm
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG2));
+      cute::copy(tScalesScale, scale_o);
+      attention_updater.rescale_o(tOrO, scale_o);
+    }
+  }
+  return;
+}
+
+template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQ,
+          typename PipelineState, typename PipelineStateQ, typename SharedStorage, typename FrgTensorO, typename AttentionUpdater>
+CUTLASS_DEVICE void mma_f16_two_stages(const Params& mainloop_params,
+                                       MainloopPipelineQ pipeline_q,
+                                       PipelineStateQ& smem_pipe_read_q,
+                                       MainloopPipeline pipeline_kv,
+                                       PipelineState& smem_pipe_read_kv,
+                                       FrgTensorO& tOrO, 
+                                       AttentionUpdater& attention_updater,
+                                       const int thread_idx, 
+                                       const int bid,
+                                       const int kv_len,
+                                       const int qo_len,
+                                       const int tile_idx,
+                                       const int q_tile_idx,
+                                       SharedStorage& shared_storage) {
+  using DTypeQ = typename Ktraits::DTypeQ;
+  using DTypeKV = typename Ktraits::DTypeKV;
+  using DTypeMD = typename Ktraits::DTypeO; 
+  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
+  using IdType = typename Ktraits::IdType;
+  using TileShape_QKD = typename Ktraits::TileShape_QKD;
+  static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS;
+  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
+  using SmemLayoutK = typename Ktraits::SmemLayoutK;
+  using SmemLayoutV = typename Ktraits::SmemLayoutV;
+  using SmemLayoutP = typename Ktraits::SmemLayoutP;
+  using SmemLayoutRow = typename Ktraits::SmemLayoutRow;
+  using SmemCopyAtom = typename Ktraits::SmemCopyAtom;
+  using SmemLayoutVt = typename Ktraits::SmemLayoutVt;
+  using SmemLayoutVtOneStage = typename Ktraits::SmemLayoutVtOneStage;
+  static_assert(is_rmem<FrgTensorO>::value, "O tensor must be rmem resident.");
+
+  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size);
+
+  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
+  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
+
+  const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
+
+  cutlass::FastDivmod stage_div(2);
+  int quotient, remainder;
+
+  Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+  Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+  // TODO: split rope/norope for sVt, not use s1~s4.
+  Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutVtOneStage{});
+  Tensor sVt_s2 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
+  Tensor sVt_s3 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 2 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
+  Tensor sVt_s4 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 3 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
+  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutP{});
+  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr), mainloop_params.layout_MD)(tile_idx, _);
+  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr), mainloop_params.layout_MD)(tile_idx, _);
+
+  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
+
+  typename Ktraits::TiledMmaQK tiled_mma_qk;
+  auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx);
+  auto smem_tiled_copy_P = make_tiled_copy_C(SmemCopyAtom{}, tiled_mma_qk);
+  auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(thread_idx);
+  Tensor tPsP = smem_thr_copy_P.partition_D(sPSS);
+  Tensor tScalesScale = s_scale(_, thread_idx % cutlass::NumThreadsPerWarpGroup, _);
+
+  typename Ktraits::TiledMmaPVSS tiled_mma_pv_ss;
+  auto threadMmaPVSS = tiled_mma_pv_ss.get_thread_slice(thread_idx);
+  Tensor tOrV1 = threadMmaPVSS.partition_fragment_B(sVt_s1);
+  Tensor tOrV2 = threadMmaPVSS.partition_fragment_B(sVt_s2);
+  Tensor tOrV3 = threadMmaPVSS.partition_fragment_B(sVt_s3);
+  Tensor tOrV4 = threadMmaPVSS.partition_fragment_B(sVt_s4);
+  Tensor tOrP_CS2 = threadMmaPVSS.partition_fragment_A(sPSS);
+
+  const int start_len = tile_idx * mainloop_params.chunk_size;
+  const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
+  const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
+  int kv_tile_idx = end_tile_idx;
+
+  auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
+    auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+    pipeline.consumer_wait(smem_pipe_read, barrier_token);
+  };
+
+  int warp_group_idx = cutlass::canonical_warp_group_idx();
+  if (warp_group_idx == 1) {
+    // consumer 0, compute qk
+    Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ);
+    Tensor tSrK = threadMmaQK.partition_fragment_B(sK);
+    auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; };
+    // wait q
+    consumer_wait(pipeline_q, smem_pipe_read_q);
+    Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+    // wait k
+    consumer_wait(pipeline_kv, smem_pipe_read_kv);
+    // first qk gemm
+    gemm</*init=*/true, /*wg_wait=*/0>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
+                                       tSrS);
+    // mask
+    {
+      Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
+      Tensor tScS = threadMmaQK.partition_C(cS);
+#pragma unroll
+      for (int i = 0; i < size(tSrS); ++i) {
+        int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
+        int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
+        if constexpr (!CAUSAL) {  // Just masking based on col
+          if (kv_idx >= kv_len) {
+            tSrS(i) = AttentionUpdater::fill_value;
+          }
+        } else {
+          if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) {
+            tSrS(i) = AttentionUpdater::fill_value;
+          }
+        }
+      }
+    }
+
+    Tensor scale_o = attention_updater.update</*init=*/true>(tSrS);
+    Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS));
+    // gather qk gemm res
+    stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
+    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, remainder));
+    cute::copy(scale_o, tScalesScale(_, remainder));
+    // r2s fence wgmma
+    cutlass::arch::fence_view_async_shared();
+    cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+
+    constexpr int n_masking_steps = CAUSAL ? cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) : 0;
+    --kv_tile_idx;
+    for (int masking_step = n_masking_steps; kv_tile_idx >= start_tile_idx; --masking_step, --kv_tile_idx) {
+      Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+      PipelineState smem_pipe_read_kv_cur = smem_pipe_read_kv;
+      ++smem_pipe_read_kv;
+      // wait next kv
+      consumer_wait(pipeline_kv, smem_pipe_read_kv);
+
+      // gemm next qk
+      gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
+                                          tSrS);
+      attention_updater.rescale_o(tOrO);
+      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv_cur.index());
+      // last pv gemm
+      if (smem_pipe_read_kv_cur.index() == 0) {
+        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                             tOrV1(_, _, _, _0{}), tOrO);
+      } else if (smem_pipe_read_kv_cur.index() == 1) {
+        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                             tOrV2(_, _, _, _0{}), tOrO);
+      } else if (smem_pipe_read_kv_cur.index() == 2) {
+        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                             tOrV3(_, _, _, _0{}), tOrO);
+      } else {
+        gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                             tOrV4(_, _, _, _0{}), tOrO);
+      }
+      // wait cur qk gemm
+      warpgroup_wait<1>();
+      // mask p
+      if (masking_step > 0) {
+        Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
+        Tensor tScS = threadMmaQK.partition_C(cS);
+#pragma unroll
+        for (int i = 0; i < size(tSrS); ++i) {
+          int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
+          int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
+          if constexpr (!CAUSAL) {  // Just masking based on col
+            if (kv_idx >= kv_len) {
+              tSrS(i) = AttentionUpdater::fill_value;
+            }
+          } else {
+            if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) {
+              tSrS(i) = AttentionUpdater::fill_value;
+            }
+          }
+        }
+      }
+      // update s (exp(s - m))
+      Tensor scale_o = attention_updater.update</*init=*/false>(tSrS);
+      Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS));
+      // gather qk gemm res
+      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
+      cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, remainder));
+      cute::copy(scale_o, tScalesScale(_, remainder));
+      // r2s fence wgmma
+      cutlass::arch::fence_view_async_shared();
+      // make sure tSrS r2s done
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+      // wait last pv gemm
+      warpgroup_wait<0>();
+      // release last kv
+      pipeline_kv.consumer_release(smem_pipe_read_kv_cur);
+    }
+    // release q
+    pipeline_q.consumer_release(smem_pipe_read_q);
+    ++smem_pipe_read_q;
+    // compute last pv
+    attention_updater.rescale_o(tOrO);
+    stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
+    if (smem_pipe_read_kv.index() == 0) {
+      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                           tOrV1(_, _, _, _0{}), tOrO);
+    } else if (smem_pipe_read_kv.index() == 1) {
+      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                           tOrV2(_, _, _, _0{}), tOrO);
+    } else if (smem_pipe_read_kv.index() == 2) {
+      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                           tOrV3(_, _, _, _0{}), tOrO);
+    } else {
+      gemm</*init=*/false, /*wg_wait=*/-1>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                           tOrV4(_, _, _, _0{}), tOrO);
+    }
+    scale_o = attention_updater.finalize(tSrS);
+    warpgroup_wait<0>();
+    // release last kv
+    pipeline_kv.consumer_release(smem_pipe_read_kv);
+    ++smem_pipe_read_kv;
+    if (chunk_num_this_seq == 1) {
+      // norm
+      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
+      cute::copy(scale_o, tScalesScale(_, remainder));
+      cutlass::arch::NamedBarrier::arrive(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2LastSync));
+      attention_updater.rescale_o(tOrO);
+    }
+    // WG1 write m,d back to gmem
+    if (chunk_num_this_seq > 1 && thread_idx % 4 == 0) { // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
+      const int warp_idx = thread_idx / 32;
+#pragma unroll
+      for (int w_i = 0; w_i < 2; ++w_i) {
+        const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i + q_group_offset;
+        const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
+
+        if (token_idx < qo_len) {
+          // const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
+          const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE;
+          const int write_idx = bid * bid_offset + token_group_idx;
+          mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
+          mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
+        }
+      }
+    }
+  } else if (warp_group_idx == 2) {
+    // consumer 1, compute pv
+    Tensor scale_o = make_tensor<DTypeQKAccum>(Shape<_2>{});
+    for (; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
+      consumer_wait(pipeline_kv, smem_pipe_read_kv);
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWarpSchedulerWG1));
+      // A: tPsP
+      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
+      cute::copy(tScalesScale(_, remainder), scale_o);
+      // rescale
+      attention_updater.rescale_o(tOrO, scale_o);
+      if (smem_pipe_read_kv.index() == 0) {
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                            tOrV1(_, _, _, _0{}), tOrO);
+      } else if (smem_pipe_read_kv.index() == 1) {
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                            tOrV2(_, _, _, _0{}), tOrO);
+      } else if (smem_pipe_read_kv.index() == 2) {
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                            tOrV3(_, _, _, _0{}), tOrO);
+      } else {
+        gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_CS2(_, _, _, remainder),
+                                            tOrV4(_, _, _, _0{}), tOrO);
+      }
+      pipeline_kv.consumer_release(smem_pipe_read_kv);
+      ++smem_pipe_read_kv;
+    }
+    if (chunk_num_this_seq == 1) {
+      // norm
+      stage_div.fast_divmod(quotient, remainder, smem_pipe_read_kv.index());
+      cutlass::arch::NamedBarrier::sync(Ktraits::NUM_MMA_THREADS, static_cast<int>(NamedBarriers::kWG1WG2LastSync));
+      cute::copy(tScalesScale(_, remainder), scale_o);
+      attention_updater.rescale_o(tOrO, scale_o);
+    }
+  }
+  return;
+}
+
+template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQ, typename MainloopPipelineQK,
+          typename PipelineState, typename PipelineStateQ, typename PipelineStateQK, typename SharedStorage, typename AttentionUpdater>
+CUTLASS_DEVICE void mma_qk_one_stages(const Params& mainloop_params,
+                                      MainloopPipelineQ pipeline_q,
+                                      PipelineStateQ& smem_pipe_read_q,
+                                      MainloopPipelineQK pipeline_qk,
+                                      PipelineStateQK& smem_pipe_write_qk,
+                                      MainloopPipeline pipeline_kv,
+                                      PipelineState& smem_pipe_read_kv,
+                                      AttentionUpdater& attention_updater,
+                                      const int& thread_idx, 
+                                      const int& bid,
+                                      const int kv_len,
+                                      const int qo_len,
+                                      const int tile_idx,
+                                      const int q_tile_idx,
+                                      SharedStorage& shared_storage) {
+  using DTypeQ = typename Ktraits::DTypeQ;
+  using DTypeKV = typename Ktraits::DTypeKV;
+  using DTypeMD = typename Ktraits::DTypeO; 
+  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
+  using IdType = typename Ktraits::IdType;
+  using TileShape_QKD = typename Ktraits::TileShape_QKD;
+  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
+  using SmemLayoutK = typename Ktraits::SmemLayoutK;
+  using SmemLayoutV = typename Ktraits::SmemLayoutV;
+  using SmemLayoutPSS = typename Ktraits::SmemLayoutPSS;
+  using SmemLayoutRow = typename Ktraits::SmemLayoutRow;
+  using SmemCopyAtom = typename Ktraits::SmemCopyAtom;
+  using SmemLayoutVt = typename Ktraits::SmemLayoutVt;
+
+  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
+  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
+
+  const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
+
+  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size);
+
+  Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+  Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+
+  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutPSS{});
+  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
+
+  typename Ktraits::TiledMmaQK tiled_mma_qk;
+  auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx);
+  Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ);
+  Tensor tSrK = threadMmaQK.partition_fragment_B(sK);
+  auto smem_tiled_copy_P = make_tiled_copy_C(SmemCopyAtom{}, tiled_mma_qk);
+  auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(thread_idx);
+  Tensor tPsP = smem_thr_copy_P.partition_D(sPSS);
+  Tensor tScalesScale = s_scale(_, thread_idx, _);
+  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr), mainloop_params.layout_MD)(tile_idx, _); // (bsz * draft_token_num * num_head)
+  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr), mainloop_params.layout_MD)(tile_idx, _);
+
+  const int start_len = tile_idx * mainloop_params.chunk_size;
+  const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
+  const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
+  int kv_tile_idx = end_tile_idx;
+
+  auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
+    auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+    pipeline.consumer_wait(smem_pipe_read, barrier_token);
+  };
+  Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+
+  // DO QK gemm
+  auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; };
+  // wait q
+  consumer_wait(pipeline_q, smem_pipe_read_q);
+  constexpr int n_masking_steps = CAUSAL ? cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) : 0;
+  int masking_step = n_masking_steps;
+  for (; kv_tile_idx >= start_tile_idx; --masking_step, --kv_tile_idx) {
+    // Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+    consumer_wait(pipeline_kv, smem_pipe_read_kv);
+    // gemm next qk
+    gemm</*init=*/true, /*wg_wait=*/0>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
+                                        tSrS);
+    // release last kv
+    pipeline_kv.consumer_release(smem_pipe_read_kv);
+    ++smem_pipe_read_kv;
+    // mask
+    if (masking_step > 0) {
+      Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
+      Tensor tScS = threadMmaQK.partition_C(cS);
+#pragma unroll
+      for (int i = 0; i < size(tSrS); ++i) {
+        int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
+        int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
+        if constexpr (!CAUSAL) {  // Just masking based on col
+          if (kv_idx >= kv_len) {
+            tSrS(i) = AttentionUpdater::fill_value;
+          }
+        } else {
+          if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) {
+            tSrS(i) = AttentionUpdater::fill_value;
+          }
+        }
+      }
+    }
+    // update s (exp(s - m))
+    Tensor scale_o = (kv_tile_idx == end_tile_idx ? attention_updater.update<true>(tSrS) : attention_updater.update<false>(tSrS));
+    Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS));
+
+    // r2s
+    pipeline_qk.producer_acquire(smem_pipe_write_qk);
+    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
+    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
+    // r2s fence wgmma
+    cutlass::arch::fence_view_async_shared();
+    pipeline_qk.producer_commit(smem_pipe_write_qk);
+    ++smem_pipe_write_qk;
+  }
+  // wg1 release q
+  pipeline_q.consumer_release(smem_pipe_read_q);
+  ++smem_pipe_read_q;
+
+  // update last softmax scales
+  Tensor scale_o = attention_updater.finalize(tSrS);
+  if (chunk_num_this_seq == 1) {
+    // norm
+    pipeline_qk.producer_acquire(smem_pipe_write_qk);
+    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
+    pipeline_qk.producer_commit(smem_pipe_write_qk);
+    ++smem_pipe_write_qk;
+  }
+
+  // WG1 write m,d back to gmem
+  if (chunk_num_this_seq > 1 && thread_idx % 4 == 0) { // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
+    const int warp_idx = thread_idx / 32;
+#pragma unroll
+    for (int w_i = 0; w_i < 2; ++w_i) {
+      const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i + q_group_offset;
+      const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
+      if (token_idx < qo_len) {
+        // const int head_idx = token_group_idx % Ktraits::GROUP_SIZE;
+        const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE;
+        const int write_idx = bid * bid_offset + token_group_idx;
+        mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
+        mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
+      }
+    }
+  }
+  return;
+}
+
+template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQ, typename MainloopPipelineQK,
+          typename PipelineState, typename PipelineStateQ, typename PipelineStateQK, typename SharedStorage, typename AttentionUpdater>
+CUTLASS_DEVICE void mma_qk_two_stages(const Params& mainloop_params,
+                                      MainloopPipelineQ pipeline_q,
+                                      PipelineStateQ& smem_pipe_read_q,
+                                      MainloopPipelineQK pipeline_qk,
+                                      PipelineStateQK& smem_pipe_write_qk,
+                                      MainloopPipeline pipeline_kv,
+                                      PipelineState& smem_pipe_read_kv,
+                                      AttentionUpdater& attention_updater,
+                                      const int& thread_idx, 
+                                      const int& bid,
+                                      const int kv_len,
+                                      const int qo_len,
+                                      const int tile_idx,
+                                      const int q_tile_idx,
+                                      SharedStorage& shared_storage) {
+  using DTypeQ = typename Ktraits::DTypeQ;
+  using DTypeKV = typename Ktraits::DTypeKV;
+  using DTypeMD = typename Ktraits::DTypeO; 
+  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
+  using IdType = typename Ktraits::IdType;
+  using TileShape_QKD = typename Ktraits::TileShape_QKD;
+  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
+  using SmemLayoutK = typename Ktraits::SmemLayoutK;
+  using SmemLayoutV = typename Ktraits::SmemLayoutV;
+  using SmemLayoutPSS = typename Ktraits::SmemLayoutPSS;
+  using SmemLayoutRow = typename Ktraits::SmemLayoutRow;
+  using SmemCopyAtom = typename Ktraits::SmemCopyAtom;
+  using SmemLayoutVt = typename Ktraits::SmemLayoutVt;
+
+  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
+  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
+
+  const int q_group_offset = q_tile_idx * BLOCK_SHAPE_Q;
+
+  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size);
+
+  Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), SmemLayoutQ{});
+  Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutK{});
+
+  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutPSS{});
+  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
+
+  typename Ktraits::TiledMmaQK tiled_mma_qk;
+  auto threadMmaQK = tiled_mma_qk.get_thread_slice(thread_idx);
+  Tensor tSrQ = threadMmaQK.partition_fragment_A(sQ);
+  Tensor tSrK = threadMmaQK.partition_fragment_B(sK);
+  auto smem_tiled_copy_P = make_tiled_copy_C(SmemCopyAtom{}, tiled_mma_qk);
+  auto smem_thr_copy_P = smem_tiled_copy_P.get_thread_slice(thread_idx);
+  Tensor tPsP = smem_thr_copy_P.partition_D(sPSS);
+  Tensor tScalesScale = s_scale(_, thread_idx, _);
+  Tensor mM = make_tensor(make_gmem_ptr(mainloop_params.m_ptr), mainloop_params.layout_MD)(tile_idx, _); // (bsz * draft_token_num * num_head)
+  Tensor mD = make_tensor(make_gmem_ptr(mainloop_params.d_ptr), mainloop_params.layout_MD)(tile_idx, _);
+
+  const int start_len = tile_idx * mainloop_params.chunk_size;
+  const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
+  int kv_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
+
+  auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
+    auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+    pipeline.consumer_wait(smem_pipe_read, barrier_token);
+  };
+
+  Tensor tSrS1 = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+  Tensor tSrS2 = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+
+  // DO QK gemm
+  auto col_limit_right = [&](int qo_idx) { return qo_idx + 1 + kv_len - qo_len; };
+  // wait q
+  consumer_wait(pipeline_q, smem_pipe_read_q);
+  // wait k
+  consumer_wait(pipeline_kv, smem_pipe_read_kv);
+  // first qk gemm
+  if (smem_pipe_write_qk.index() == 0) {
+    gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
+                                        tSrS1);
+  } else {
+    gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
+                                        tSrS2);
+  }
+  constexpr int n_masking_steps = CAUSAL ? cute::ceil_div(BLOCK_SHAPE_Q, BLOCK_SHAPE_KV) : 0;
+  int masking_step = n_masking_steps;
+  --kv_tile_idx;
+  for (; kv_tile_idx >= start_tile_idx; --masking_step, --kv_tile_idx) {
+    Tensor tSrS = partition_fragment_C(tiled_mma_qk, select<0, 1>(TileShape_QKD{}));
+    PipelineState smem_pipe_read_kv_cur = smem_pipe_read_kv;
+    ++smem_pipe_read_kv;
+    // wait next kv
+    consumer_wait(pipeline_kv, smem_pipe_read_kv);
+    // gemm next qk
+    if (smem_pipe_write_qk.index() == 1) {
+      gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
+                                          tSrS1);
+    } else {
+      gemm</*init=*/true, /*wg_wait=*/-1>(tiled_mma_qk, tSrQ, tSrK(_, _, _, smem_pipe_read_kv.index()),
+                                          tSrS2);
+    }
+    // wait last qk done
+    warpgroup_wait<1>();
+    // release last kv
+    pipeline_kv.consumer_release(smem_pipe_read_kv_cur);
+    // mask now
+    if (masking_step > 0) {
+      Tensor cS = cute::make_identity_tensor(select<0, 1>(TileShape_QKD{}));
+      Tensor tScS = threadMmaQK.partition_C(cS);
+#pragma unroll
+      for (int i = 0; i < size(tSrS1); ++i) {
+        int qo_idx = (get<0>(tScS(i)) + q_group_offset) / Ktraits::GROUP_SIZE;
+        int kv_idx = get<1>(tScS(i)) + kv_tile_idx * BLOCK_SHAPE_KV;
+        if (kv_idx >= std::min(kv_len, col_limit_right(qo_idx))) {
+          if (smem_pipe_write_qk.index() == 0) {
+            tSrS1(i) = AttentionUpdater::fill_value;
+          } else {
+            tSrS2(i) = AttentionUpdater::fill_value;
+          }
+        }
+      }
+    }
+    // update s (exp(s - m))
+    if (smem_pipe_write_qk.index() == 0) {
+      Tensor scale_o = attention_updater.update</*init=*/false>(tSrS1);
+      Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS1));
+      // r2s
+      pipeline_qk.producer_acquire(smem_pipe_write_qk);
+      cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
+      cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
+    } else {
+      Tensor scale_o = attention_updater.update</*init=*/false>(tSrS2);
+      Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS2));
+      // r2s
+      pipeline_qk.producer_acquire(smem_pipe_write_qk);
+      cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
+      cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
+    }
+    // r2s fence wgmma
+    cutlass::arch::fence_view_async_shared();
+    pipeline_qk.producer_commit(smem_pipe_write_qk);
+    ++smem_pipe_write_qk;
+  }
+  // wg1 release q
+  pipeline_q.consumer_release(smem_pipe_read_q);
+  ++smem_pipe_read_q;
+  // wait last qk
+  warpgroup_wait<0>();
+  // wg1 release kv
+  pipeline_kv.consumer_release(smem_pipe_read_kv);
+  ++smem_pipe_read_kv;
+  // update s (exp(s - m)) last
+  if (smem_pipe_write_qk.index() == 0) {
+    Tensor scale_o = attention_updater.update</*init=*/false>(tSrS1);
+    Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS1));
+    // r2s
+    pipeline_qk.producer_acquire(smem_pipe_write_qk);
+    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
+    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
+  } else {
+    Tensor scale_o = attention_updater.update</*init=*/false>(tSrS2);
+    Tensor tPrP = smem_thr_copy_P.retile_S(convert_type<DTypeKV>(tSrS2));
+    // r2s
+    pipeline_qk.producer_acquire(smem_pipe_write_qk);
+    cute::copy(smem_tiled_copy_P, tPrP, tPsP(_, _, _, smem_pipe_write_qk.index()));
+    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
+  }
+  // r2s fence wgmma
+  cutlass::arch::fence_view_async_shared();
+  pipeline_qk.producer_commit(smem_pipe_write_qk);
+  ++smem_pipe_write_qk;
+
+  // update last softmax scales
+  Tensor scale_o = attention_updater.finalize(tSrS1);
+  if (chunk_num_this_seq == 1) {
+    // norm
+    pipeline_qk.producer_acquire(smem_pipe_write_qk);
+    cute::copy(scale_o, tScalesScale(_, smem_pipe_write_qk.index()));
+    pipeline_qk.producer_commit(smem_pipe_write_qk);
+    ++smem_pipe_write_qk;
+  }
+
+  // WG1 write m,d back to gmem
+  if (chunk_num_this_seq > 1 && thread_idx % 4 == 0) { // 16 rows per warp, eg. t0->row0 row8，t4->row1 row9
+    const int warp_idx = thread_idx / 32;
+#pragma unroll
+    for (int w_i = 0; w_i < 2; ++w_i) {
+      const int token_group_idx = warp_idx * 16 + (thread_idx % 32) / 4 + 8 * w_i + q_group_offset;
+      const int token_idx = token_group_idx / Ktraits::GROUP_SIZE;
+      if (token_idx < qo_len) {
+        const int bid_offset = mainloop_params.draft_total_token_num * Ktraits::GROUP_SIZE;
+        const int write_idx = bid * bid_offset + token_group_idx;
+        mM(write_idx) = static_cast<DTypeMD>(attention_updater.row_max(w_i));
+        mD(write_idx) = static_cast<DTypeMD>(attention_updater.row_sum(w_i));
+      }
+    }
+  }
+  return;
+}
+
+template <typename Ktraits, bool CAUSAL, typename Params, typename MainloopPipeline, typename MainloopPipelineQK,
+          typename PipelineState, typename PipelineStateQK, typename SharedStorage, typename FrgTensorO, typename AttentionUpdater>
+CUTLASS_DEVICE void mma_pv_one_stages(const Params& mainloop_params,
+                                      MainloopPipelineQK pipeline_qk,
+                                      PipelineStateQK& smem_pipe_read_qk,
+                                      MainloopPipeline pipeline_kv,
+                                      PipelineState& smem_pipe_read_kv,
+                                      AttentionUpdater& attention_updater,
+                                      FrgTensorO& tOrO, 
+                                      const int thread_idx, 
+                                      const int bid,
+                                      const int kv_len,
+                                      const int qo_len,
+                                      const int tile_idx,
+                                      SharedStorage& shared_storage) {
+  using DTypeQ = typename Ktraits::DTypeQ;
+  using DTypeKV = typename Ktraits::DTypeKV;
+  using DTypeMD = typename Ktraits::DTypeO;
+  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
+  using IdType = typename Ktraits::IdType;
+  using TileShape_QKD = typename Ktraits::TileShape_QKD;
+  using SmemLayoutQ = typename Ktraits::SmemLayoutQ;
+  using SmemLayoutK = typename Ktraits::SmemLayoutK;
+  using SmemLayoutV = typename Ktraits::SmemLayoutV;
+  using SmemLayoutPSS = typename Ktraits::SmemLayoutPSS;
+  using SmemLayoutRow = typename Ktraits::SmemLayoutRow;
+  using SmemCopyAtom = typename Ktraits::SmemCopyAtom;
+  using SmemLayoutVtOneStage = typename Ktraits::SmemLayoutVtOneStage;
+  static_assert(is_rmem<FrgTensorO>::value, "O tensor must be rmem resident.");
+
+  static constexpr int BLOCK_SHAPE_Q = get<0>(TileShape_QKD{});
+  static constexpr int BLOCK_SHAPE_KV = get<1>(TileShape_QKD{});
+
+  const int chunk_num_this_seq = cute::ceil_div(kv_len, mainloop_params.chunk_size);
+  // todo: split rope/norope for sVt, not use s1~s4.
+  Tensor sVt_s1 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data()), SmemLayoutVtOneStage{});
+  Tensor sVt_s2 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
+  Tensor sVt_s3 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 2 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
+  Tensor sVt_s4 = make_tensor(make_smem_ptr(shared_storage.smem_kv.data() + 3 * Ktraits::NUM_PER_STAGE), SmemLayoutVtOneStage{});
+
+  Tensor sPSS = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), SmemLayoutPSS{});
+  Tensor s_scale = make_tensor(make_smem_ptr(shared_storage.smem_scale.data()), SmemLayoutRow{});
+  Tensor tScalesScale = s_scale(_, thread_idx % cutlass::NumThreadsPerWarpGroup, _);
+
+  typename Ktraits::TiledMmaPVSS tiled_mma_pv_ss;
+  auto threadMmaPVSS = tiled_mma_pv_ss.get_thread_slice(thread_idx);
+  Tensor tOrV1 = threadMmaPVSS.partition_fragment_B(sVt_s1);
+  Tensor tOrV2 = threadMmaPVSS.partition_fragment_B(sVt_s2);
+  Tensor tOrV3 = threadMmaPVSS.partition_fragment_B(sVt_s3);
+  Tensor tOrV4 = threadMmaPVSS.partition_fragment_B(sVt_s4);
+  Tensor tOrP_S = threadMmaPVSS.partition_fragment_A(sPSS);
+  const int start_len = tile_idx * mainloop_params.chunk_size;
+  const int start_tile_idx = start_len / BLOCK_SHAPE_KV;
+  const int end_tile_idx = cute::ceil_div(min(start_len + mainloop_params.chunk_size, kv_len), BLOCK_SHAPE_KV) - 1;
+  int kv_tile_idx = end_tile_idx;
+
+  auto consumer_wait = [](auto& pipeline, auto& smem_pipe_read) {
+    auto barrier_token = pipeline.consumer_try_wait(smem_pipe_read);
+    pipeline.consumer_wait(smem_pipe_read, barrier_token);
+  };
+
+  Tensor scale_o = make_tensor<DTypeQKAccum>(Shape<_2>{});
+  for (; kv_tile_idx >= start_tile_idx; --kv_tile_idx) {
+    // wait v
+    consumer_wait(pipeline_kv, smem_pipe_read_kv);
+    // wait p
+    consumer_wait(pipeline_qk, smem_pipe_read_qk);
+    // A: tPsP
+    cute::copy(tScalesScale(_, smem_pipe_read_qk.index()), scale_o);
+    // rescale
+    attention_updater.rescale_o(tOrO, scale_o);
+    // gemm pv
+    if (smem_pipe_read_kv.index() == 0) {
+      gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_S(_, _, _, smem_pipe_read_qk.index()),
+                                          tOrV1(_, _, _, _0{}), tOrO);
+    } else if (smem_pipe_read_kv.index() == 1) {
+      gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_S(_, _, _, smem_pipe_read_qk.index()),
+                                          tOrV2(_, _, _, _0{}), tOrO);
+    } else if (smem_pipe_read_kv.index() == 2) {
+      gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_S(_, _, _, smem_pipe_read_qk.index()),
+                                          tOrV3(_, _, _, _0{}), tOrO);
+    } else if (smem_pipe_read_kv.index() == 3) {
+      gemm</*init=*/false, /*wg_wait=*/0>(tiled_mma_pv_ss, tOrP_S(_, _, _, smem_pipe_read_qk.index()),
+                                          tOrV4(_, _, _, _0{}), tOrO);
+    }
+    pipeline_qk.consumer_release(smem_pipe_read_qk);
+    ++smem_pipe_read_qk;
+    pipeline_kv.consumer_release(smem_pipe_read_kv);
+    ++smem_pipe_read_kv;
+  }
+  if (chunk_num_this_seq == 1) {
+    // norm
+    consumer_wait(pipeline_qk, smem_pipe_read_qk);
+    cute::copy(tScalesScale(_, smem_pipe_read_qk.index()), scale_o);
+    pipeline_qk.consumer_release(smem_pipe_read_qk);
+    ++smem_pipe_read_qk;
+    attention_updater.rescale_o(tOrO, scale_o);
+  }
+  return;
+}
+
+}  // namespace mla_attn
+
+#endif  // ATTENTION_HOPPER_MAINLOOP_MMA_CUH_

csrc/gpu/mla_attn/mla_hopper.cuh

@@ -0,0 +1,484 @@
+// 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.
+
+/*
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
+ * Dao. Licensed under the BSD 3-Clause.
+ *
+ * Modified by the FlashInfer team.
+ */
+
+#pragma once
+
+#ifndef ATTENTION_HOPPER_PREFILL_SM90_CUH_
+#define ATTENTION_HOPPER_PREFILL_SM90_CUH_
+
+#include <cuda.h>
+#include <cuda_device_runtime_api.h>
+#include <cutlass/arch/reg_reconfig.h>
+#include <cutlass/array.h>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/numeric_types.h>
+
+#include <type_traits>
+#include <vector>
+
+#include "attention_updater.cuh"
+#include "cute/tensor.hpp"
+#include "cutlass/pipeline/pipeline.hpp"
+#include "epilogue.cuh"
+#include "helper.h"
+#include "kernel_traits.cuh"
+#include "mainloop_mma.cuh"
+#include "mainloop_load.cuh"
+#include "utils.cuh"
+
+#ifdef DEBUG_MLA
+#undef DEBUG_MLA
+#endif
+// #define DEBUG_MLA
+
+namespace mla_attn {
+
+using namespace cute;
+
+template <typename DTypeQ_, typename DTypeKV_, typename DTypeO_, typename IdType_>
+struct Params {
+    using DTypeQ = DTypeQ_;
+    using DTypeKV = DTypeKV_;
+    using DTypeO = DTypeO_;
+    using IdType = IdType_;
+
+    alignas(16) DTypeQ *Q; // [token_num, head_num, dim_head]
+    alignas(16) DTypeKV *KV; // [max_block_num, block_size, dim_head]
+    alignas(16) DTypeO *O; // [token_num, head_num, dim_head]
+    alignas(16) DTypeO *O_tmp; // [num_chunks, bsz, head_num, dim_head]
+    alignas(16) float *m; // [num_chunks, bsz * draft_total_token_num * head_num]
+    alignas(16) float *d; // [num_chunks, bsz * draft_total_token_num * head_num]
+
+    alignas(16) IdType *block_tables;
+    alignas(16) IdType *seq_lens_this_time;
+    alignas(16) IdType *seq_lens_encoder;
+    alignas(16) IdType *seq_lens_decoder;
+    alignas(16) IdType *cumsum_q_seqlens;
+    alignas(16) IdType *padding_offsets;
+
+    alignas(16) IdType *batch_ids;
+    alignas(16) IdType *q_tile_ids_per_batch = nullptr;
+    alignas(16) IdType *tile_ids_per_batch = nullptr;
+    alignas(16) IdType *num_blocks_x;
+
+
+    uint32_t q_stride_bsz;
+    uint32_t q_stride_head_num;
+
+    uint32_t kv_stride_block_num;
+    uint32_t kv_stride_block_size;
+
+    uint32_t o_stride_bsz;
+    uint32_t o_stride_head_num;
+
+    int bsz;
+    int token_num;
+    int max_seq_len;
+    int max_block_num;
+    int max_block_num_per_seq;
+    int q_num_head;
+    int qk_head_dim;
+    int vo_head_dim;
+    int block_size;
+    int draft_total_token_num;
+    int chunk_size;
+    int chunk_num;
+    int num_blocks_x_int;
+
+    float sm_scale;
+};
+
+#define DISPATCH_GROUP_SIZE(group_size, GROUP_SIZE, ...)     \
+  if (group_size == 8) {                                     \
+    constexpr size_t GROUP_SIZE = 8;                         \
+    __VA_ARGS__                                              \
+  } else if (group_size == 16) {                             \
+    constexpr size_t GROUP_SIZE = 16;                        \
+    __VA_ARGS__                                              \
+  } else if (group_size == 64) {                             \
+    constexpr size_t GROUP_SIZE = 64;                        \
+    __VA_ARGS__                                              \
+  } else if (group_size == 128) {                            \
+    constexpr size_t GROUP_SIZE = 128;                       \
+    __VA_ARGS__                                              \
+  } else {                                                   \
+    PD_THROW("not support the group_size: ", group_size);    \
+    return cudaErrorNotSupported;                            \
+  }
+template <typename CollectiveMainloop, typename CollectiveEpilogue, typename Ktraits, bool CAUSAL, bool USE_REG_EALLOC = false>
+__global__ void __launch_bounds__(Ktraits::NUM_WARPS * cutlass::NumThreadsPerWarp, 1)
+MLAWithKVCacheKernel(CUTE_GRID_CONSTANT
+                     typename CollectiveMainloop::Params const mainloop_params,
+                     CUTE_GRID_CONSTANT
+                     typename CollectiveEpilogue::Params const epilogue_params) {
+
+  using DTypeQ = typename Ktraits::DTypeQ;
+  using DTypeKV = typename Ktraits::DTypeKV;
+  using DTypeO = typename Ktraits::DTypeO;
+  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
+  using TileShape_QKD = typename Ktraits::TileShape_QKD;
+  using TileShape_PDV = typename Ktraits::TileShape_PDV;
+
+  static constexpr int NUM_MMA_THREADS = Ktraits::NUM_MMA_THREADS;
+  static constexpr int NUM_COPY_THREADS = Ktraits::NUM_PRODUCER_THREADS;
+  static constexpr int BLOCK_SHAPE_Q = Ktraits::BLOCK_SHAPE_Q;
+  static constexpr int BLOCK_SHAPE_KV = Ktraits::BLOCK_SHAPE_KV;
+  const int num_blocks_x = mainloop_params.num_blocks_x[0];
+
+  static constexpr bool use_tma_load_kv = CollectiveMainloop::USE_TMA_LOAD_KV;
+
+  using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState = typename MainloopPipeline::PipelineState;
+
+  using MainloopPipelineQ = typename CollectiveMainloop::MainloopPipelineQ;
+  using PipelineParamsQ = typename MainloopPipelineQ::Params;
+  using PipelineStateQ = typename MainloopPipelineQ::PipelineState;
+
+  extern __shared__ char shared_memory[];
+  auto& shared_storage = *reinterpret_cast<typename Ktraits::SharedStorage*>(shared_memory);
+
+  int const lane_predicate = cute::elect_one_sync();
+  int const warp_idx = cutlass::canonical_warp_idx_sync();
+
+  if (warp_idx == 0 && lane_predicate) {
+    CollectiveMainloop::prefetch_tma_descriptors(mainloop_params);
+    CollectiveEpilogue::prefetch_tma_descriptors(epilogue_params);
+  }
+
+  // Obtain warp index
+  int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup;
+
+  PipelineParams pipeline_params;
+  int warp_group_idx = cutlass::canonical_warp_group_idx();
+  pipeline_params.role = warp_group_idx == 0 ? MainloopPipeline::ThreadCategory::Producer
+                                             : MainloopPipeline::ThreadCategory::Consumer;
+  if constexpr (use_tma_load_kv) {
+    pipeline_params.is_leader = warp_group_thread_idx == 0;
+    pipeline_params.num_consumers = NUM_MMA_THREADS;
+  } else {
+    pipeline_params.producer_arv_count = NUM_COPY_THREADS;
+    pipeline_params.consumer_arv_count = NUM_MMA_THREADS;
+  }
+
+  PipelineParamsQ pipeline_params_q;
+  pipeline_params_q.role = warp_group_idx == 0 ? MainloopPipelineQ::ThreadCategory::Producer
+                                               : MainloopPipelineQ::ThreadCategory::Consumer;
+  pipeline_params_q.producer_arv_count = NUM_COPY_THREADS;
+  pipeline_params_q.consumer_arv_count = cutlass::NumThreadsPerWarpGroup; // just one wg qk
+  
+
+  MainloopPipelineQ pipeline_q(shared_storage.pipeline_q, pipeline_params_q);
+  MainloopPipeline pipeline_kv = [&] {
+    if constexpr (use_tma_load_kv) {
+      pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesKV;
+      return MainloopPipeline(shared_storage.pipeline_kv, pipeline_params,
+                              /*cluster_shape=*/Shape<_1, _1, _1>{});
+    } else {
+      return MainloopPipeline(shared_storage.pipeline_kv, pipeline_params);
+    }
+  }();
+  __syncthreads();
+
+  CollectiveMainloop collective_mainloop;
+  CollectiveEpilogue collective_epilogue;
+  
+  if (warp_group_idx == 0) {
+    // producer
+    if (USE_REG_EALLOC) {
+      cutlass::arch::warpgroup_reg_dealloc<88>();
+    }
+    const uint32_t warp_idx_in_warpgroup = __shfl_sync(0xffffffff, warp_idx % 4, 0);
+    
+    PipelineStateQ smem_pipe_write_q = cutlass::make_producer_start_state<MainloopPipelineQ>();
+    PipelineState smem_pipe_write_kv = cutlass::make_producer_start_state<MainloopPipeline>();
+    const int block_id = blockIdx.x;
+    const int bid = mainloop_params.batch_ids[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch ? mainloop_params.tile_ids_per_batch[block_id] : 0;
+    const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
+    const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+    const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
+    const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+    const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+
+    // load Q
+    collective_mainloop.load_q(
+        mainloop_params,
+        pipeline_q,
+        smem_pipe_write_q,
+        shared_storage,
+        threadIdx.x,
+        bid,
+        q_tile_id);
+
+    if constexpr (!use_tma_load_kv) {
+      // load kv
+      collective_mainloop.load_kv(
+          mainloop_params,
+          pipeline_kv,
+          smem_pipe_write_kv,
+          shared_storage,
+          bid,
+          seq_len_decoder_now,
+          tile_id
+      );
+    } else {
+      if (warp_idx_in_warpgroup == 0) {
+        // load kv tma
+        collective_mainloop.load_kv_tma(
+            mainloop_params,
+            pipeline_kv,
+            smem_pipe_write_kv,
+            shared_storage,
+            bid,
+            seq_len_decoder_now,
+            tile_id
+        );
+      }
+    }
+  } else {
+    // consumer
+    if (USE_REG_EALLOC) {
+      cutlass::arch::warpgroup_reg_alloc<208>();
+    }
+    PipelineStateQ smem_pipe_read_q;
+    PipelineState smem_pipe_read_kv;
+
+    typename Ktraits::TiledMmaPVSS tiled_mma_pv;
+    Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_PDV{}));
+
+    auto attention_updater = OnlineSoftmax<2 * size<1>(tOrO), /*WITH_SCALE=*/true>(mainloop_params.sm_scale);
+    const int block_id = blockIdx.x;
+    clear(tOrO);
+    clear(attention_updater.scores_scale);
+    const int bid = mainloop_params.batch_ids[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch ? mainloop_params.tile_ids_per_batch[block_id] : 0;
+    const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
+    const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+    const int seq_len_encoder_now = mainloop_params.seq_lens_encoder[bid];
+    const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+    const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+
+    if constexpr (BLOCK_SHAPE_KV == 64) {
+      mma_f16<Ktraits, CAUSAL>(
+        mainloop_params, 
+        pipeline_q, 
+        smem_pipe_read_q,
+        pipeline_kv, 
+        smem_pipe_read_kv,
+        tOrO, 
+        attention_updater,
+        threadIdx.x - NUM_COPY_THREADS,
+        bid,
+        seq_len_decoder_now,
+        seq_len_now,
+        tile_id,
+        q_tile_id,
+        shared_storage);
+    } else if (BLOCK_SHAPE_KV == 32) {
+      mma_f16_two_stages<Ktraits, CAUSAL>(
+        mainloop_params, 
+        pipeline_q, 
+        smem_pipe_read_q,
+        pipeline_kv, 
+        smem_pipe_read_kv,
+        tOrO, 
+        attention_updater,
+        threadIdx.x - NUM_COPY_THREADS,
+        bid,
+        seq_len_decoder_now,
+        seq_len_now,
+        tile_id,
+        q_tile_id,
+        shared_storage);
+    }
+
+    collective_epilogue.store(
+        epilogue_params, 
+        tOrO, 
+        attention_updater.get_lse(),
+        shared_storage,
+        tiled_mma_pv, 
+        threadIdx.x - NUM_COPY_THREADS,
+        bid,
+        mainloop_params.bsz,
+        seq_len_now,
+        start_token_idx,
+        tile_id,
+        q_tile_id,
+        seq_len_decoder_now,
+        mainloop_params.chunk_size,
+        mainloop_params.draft_total_token_num,
+        mainloop_params.o_stride_head_num);
+  }
+}
+
+
+template <typename KernelTraits, bool CAUSAL, typename Params, bool USE_REG_EALLOC=false>
+cudaError_t BatchMLAWithPagedKVCacheKernelTraitsDispatched(Params& params,
+                                                           cudaStream_t stream) {
+  using DTypeQ = typename KernelTraits::DTypeQ;
+  using DTypeKV = typename KernelTraits::DTypeKV;
+  using DTypeO = typename KernelTraits::DTypeO;
+  using IdType = typename KernelTraits::IdType;
+  using NV_TYPE = typename KernelTraits::NV_TYPE;
+
+  using CollectiveMainloop =
+      CollectiveMainloop<KernelTraits, CAUSAL>;
+  using CollectiveEpilogue = CollectiveEpilogue<KernelTraits>;
+
+  typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments({
+      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.qk_head_dim), make_stride(params.qk_head_dim, _1{})), // layout q
+      make_layout(make_shape(params.block_size, params.qk_head_dim, params.max_block_num), make_stride(params.qk_head_dim, _1{}, params.block_size * params.qk_head_dim)),
+      make_layout(make_shape(params.chunk_num, params.bsz * params.draft_total_token_num * params.q_num_head), make_stride(params.bsz * params.draft_total_token_num * params.q_num_head, _1{})),
+      params.Q,
+      params.KV,
+      params.m,
+      params.d,
+      params.block_tables,
+      params.seq_lens_this_time,
+      params.seq_lens_encoder,
+      params.seq_lens_decoder,
+      params.cumsum_q_seqlens,
+      params.batch_ids,
+      params.tile_ids_per_batch,
+      params.q_tile_ids_per_batch,
+      params.num_blocks_x,
+      params.sm_scale,
+      params.bsz,
+      params.max_block_num,
+      params.max_block_num_per_seq,
+      params.q_stride_bsz,
+      params.q_stride_head_num,
+      params.kv_stride_block_num,
+      params.kv_stride_block_size,
+      params.o_stride_bsz,
+      params.o_stride_head_num,
+      params.chunk_size,
+      params.chunk_num,
+      params.draft_total_token_num
+  });
+  typename CollectiveEpilogue::Params epilogue_params = CollectiveEpilogue::to_underlying_arguments_ntma({
+      params.O,
+      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.vo_head_dim), make_stride(params.vo_head_dim, _1{})), // layout O
+      params.O_tmp,
+      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.vo_head_dim), make_stride(params.vo_head_dim, _1{})) // layout O_tmp
+  });
+
+  // Get the ptr to kernel function.
+  auto kernel =
+      MLAWithKVCacheKernel<CollectiveMainloop, CollectiveEpilogue, KernelTraits, CAUSAL, USE_REG_EALLOC>;
+  int smem_size = sizeof(typename KernelTraits::SharedStorage);
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  int device;
+  cudaGetDevice(&device);
+  int multiprocessor_count;
+  cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device);
+  int act_blocks_per_sm;
+  cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+      &act_blocks_per_sm, kernel, KernelTraits::NUM_WARPS * 32, smem_size);
+
+  dim3 grid_dims = {params.num_blocks_x_int, 1, 1}; // todo: split kv
+  static constexpr int ctaSize = KernelTraits::NUM_WARPS * 32;
+  dim3 block_dims(ctaSize, 1, 1);
+  kernel<<<grid_dims, block_dims, smem_size, stream>>>(
+    mainloop_params, epilogue_params
+  );
+  if (params.chunk_num > 1) {
+    constexpr int vec_size = 16 / sizeof(DTypeO);
+    constexpr int merge_block_size = 256;
+    constexpr int blockx = KernelTraits::HEAD_DIM_VO / vec_size;
+    constexpr int blocky = (merge_block_size + blockx - 1) / blockx;
+    dim3 grids_merge(min(multiprocessor_count, params.token_num), params.q_num_head); // 128k is too large
+    dim3 blocks_merge(blockx, blocky);
+    merge_multi_chunks_kernel<NV_TYPE, vec_size, blocky, KernelTraits::HEAD_DIM_VO><<<grids_merge, blocks_merge, 0, stream>>>(
+      reinterpret_cast<NV_TYPE*>(params.O_tmp),
+      params.m,
+      params.d,
+      params.seq_lens_this_time,
+      params.seq_lens_decoder,
+      params.seq_lens_encoder,
+      params.padding_offsets,
+      reinterpret_cast<NV_TYPE*>(params.O),
+      params.max_seq_len,
+      params.chunk_num,
+      params.q_num_head,
+      params.chunk_size,
+      params.vo_head_dim,
+      params.token_num,
+      params.bsz,
+      params.draft_total_token_num
+    );
+  }
+  return cudaSuccess;
+}
+
+template <uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO, typename NV_TYPE, typename Params, bool USE_REG_EALLOC=false>
+cudaError_t BatchMLAWithPagedKVCacheDispatched(Params& params, cudaStream_t stream) {
+  constexpr bool CAUSAL = true;
+  if constexpr (HEAD_DIM_QK == 576) {
+    if (params.block_size == 32) {
+      DISPATCH_GROUP_SIZE(params.q_num_head, GROUP_SIZE,
+        BatchMLAWithPagedKVCacheKernelTraitsDispatched<
+            AttentionKernelTraits</*USE_TMA_LOAD_KV=*/true, 
+                                  HEAD_DIM_QK, 
+                                  HEAD_DIM_VO, 
+                                  GROUP_SIZE,
+                                  /*BLOCK_SHAPE_Q_=*/64,
+                                  /*BLOCK_SHAPE_KV_=*/32,
+                                  /*NUM_STAGES_=*/4, 
+                                  typename Params::DTypeQ,
+                                  typename Params::DTypeKV, 
+                                  typename Params::DTypeO,
+                                  typename Params::IdType,
+                                  NV_TYPE>,
+            CAUSAL,
+            Params,
+            USE_REG_EALLOC>(params, stream);)
+    } else if (params.block_size == 64) {
+      DISPATCH_GROUP_SIZE(params.q_num_head, GROUP_SIZE,
+        BatchMLAWithPagedKVCacheKernelTraitsDispatched<
+            AttentionKernelTraits</*USE_TMA_LOAD_KV=*/true, 
+                                  HEAD_DIM_QK, 
+                                  HEAD_DIM_VO, 
+                                  GROUP_SIZE,
+                                  /*BLOCK_SHAPE_Q_=*/64,
+                                  /*BLOCK_SHAPE_KV_=*/64,
+                                  /*NUM_STAGES_=*/2, 
+                                  typename Params::DTypeQ,
+                                  typename Params::DTypeKV, 
+                                  typename Params::DTypeO,
+                                  typename Params::IdType,
+                                  NV_TYPE>,
+            CAUSAL,
+            Params,
+            USE_REG_EALLOC>(params, stream);)
+    }
+  } else {
+    return cudaErrorNotSupported;
+  }
+  return cudaSuccess;
+};
+
+}  // namespace mla_attn
+
+#endif  // ATTENTION_HOPPER_PREFILL_SM90_CUH_

csrc/gpu/mla_attn/mla_hopper_wg4.cuh

@@ -0,0 +1,426 @@
+// 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
+
+#ifndef ATTENTION_HOPPER_PREFILL_WG4_SM90_CUH_
+#define ATTENTION_HOPPER_PREFILL_WG4_SM90_CUH_
+
+#include <cuda.h>
+#include <cuda_device_runtime_api.h>
+#include <cutlass/arch/reg_reconfig.h>
+#include <cutlass/array.h>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/numeric_types.h>
+
+#include <type_traits>
+#include <vector>
+
+#include "attention_updater.cuh"
+#include "cute/tensor.hpp"
+#include "cutlass/pipeline/pipeline.hpp"
+#include "epilogue.cuh"
+#include "helper.h"
+#include "kernel_traits_wg4.cuh"
+#include "mla_hopper.cuh"
+#include "mainloop_mma.cuh"
+#include "mainloop_load.cuh"
+#include "utils.cuh"
+
+#ifdef DEBUG_MLA
+#undef DEBUG_MLA
+#endif
+// #define DEBUG_MLA
+
+namespace mla_attn {
+
+using namespace cute;
+
+template <typename CollectiveMainloop, typename CollectiveEpilogue, typename Ktraits, bool CAUSAL, bool USE_REG_EALLOC=false, bool USE_QK_TWO_STAGE=false>
+__global__ void __launch_bounds__(Ktraits::NUM_WARPS * cutlass::NumThreadsPerWarp, 1)
+MLAWithKVCacheWG4Kernel(CUTE_GRID_CONSTANT
+                        typename CollectiveMainloop::Params const mainloop_params,
+                        CUTE_GRID_CONSTANT
+                        typename CollectiveEpilogue::Params const epilogue_params) {
+  int warp_group_idx = cutlass::canonical_warp_group_idx();
+  using DTypeQ = typename Ktraits::DTypeQ;
+  using DTypeKV = typename Ktraits::DTypeKV;
+  using DTypeO = typename Ktraits::DTypeO;
+  using DTypeQKAccum = typename Ktraits::DTypeQKAccum;
+  using TileShape_QKD = typename Ktraits::TileShape_QKD;
+  using TileShape_PDV = typename Ktraits::TileShape_PDV;
+
+  static constexpr int NUM_MMA_THREADS_QK = Ktraits::NUM_MMA_THREADS_QK;
+  static constexpr int NUM_MMA_THREADS_PV = Ktraits::NUM_MMA_THREADS_PV;
+  static constexpr int NUM_COPY_THREADS = Ktraits::NUM_PRODUCER_THREADS;
+  static constexpr int BLOCK_SHAPE_Q = Ktraits::BLOCK_SHAPE_Q;
+  static constexpr int BLOCK_SHAPE_KV = Ktraits::BLOCK_SHAPE_KV;
+  const int num_blocks_x = mainloop_params.num_blocks_x[0];
+
+  static constexpr bool use_tma_load_kv = CollectiveMainloop::USE_TMA_LOAD_KV;
+
+  using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState = typename MainloopPipeline::PipelineState;
+
+  using MainloopPipelineQ = typename CollectiveMainloop::MainloopPipelineQ;
+  using PipelineParamsQ = typename MainloopPipelineQ::Params;
+  using PipelineStateQ = typename MainloopPipelineQ::PipelineState;
+
+  using MainloopPipelineQK = typename Ktraits::MainloopPipelineQK;
+  using PipelineParamsQK = typename MainloopPipelineQK::Params;
+  using PipelineStateQK = typename MainloopPipelineQK::PipelineState;
+
+  extern __shared__ char shared_memory[];
+  auto& shared_storage = *reinterpret_cast<typename Ktraits::SharedStorage*>(shared_memory);
+
+  PipelineParams pipeline_params;
+  pipeline_params.role = warp_group_idx == 0 ? MainloopPipeline::ThreadCategory::Producer
+                                             : MainloopPipeline::ThreadCategory::Consumer;
+  if constexpr (use_tma_load_kv) {
+    int const warp_group_thread_idx = threadIdx.x % cutlass::NumThreadsPerWarpGroup;
+    pipeline_params.is_leader = warp_group_thread_idx == 0;
+    pipeline_params.num_consumers = NUM_MMA_THREADS_QK + NUM_MMA_THREADS_PV;
+  } else {
+    pipeline_params.producer_arv_count = NUM_COPY_THREADS;
+    pipeline_params.consumer_arv_count = NUM_MMA_THREADS_QK + NUM_MMA_THREADS_PV;
+  }
+
+  PipelineParamsQ pipeline_params_q;
+  if (warp_group_idx == 0) {
+    pipeline_params_q.role = MainloopPipelineQ::ThreadCategory::Producer;
+  } else if (warp_group_idx == 1) {
+    pipeline_params_q.role = MainloopPipelineQ::ThreadCategory::Consumer;
+  } else {
+    pipeline_params_q.role = MainloopPipelineQ::ThreadCategory::NonParticipant;
+  }
+  pipeline_params_q.producer_arv_count = NUM_COPY_THREADS;
+  pipeline_params_q.consumer_arv_count = cutlass::NumThreadsPerWarpGroup; // just one wg qk
+  
+
+  MainloopPipelineQ pipeline_q(shared_storage.pipeline_q, pipeline_params_q);
+
+  PipelineParamsQK pipeline_params_qk;
+  if (warp_group_idx == 0) {
+    pipeline_params_qk.role = MainloopPipelineQK::ThreadCategory::NonParticipant;
+  } else if (warp_group_idx == 1) {
+    pipeline_params_qk.role = MainloopPipelineQK::ThreadCategory::Producer;
+  } else {
+    pipeline_params_qk.role = MainloopPipelineQK::ThreadCategory::Consumer;
+  }
+  pipeline_params_qk.producer_arv_count = NUM_MMA_THREADS_QK;
+  pipeline_params_qk.consumer_arv_count = NUM_MMA_THREADS_PV;
+  MainloopPipelineQK pipeline_qk(shared_storage.pipeline_qk, pipeline_params_qk);
+
+  MainloopPipeline pipeline_kv = [&] {
+    if constexpr (use_tma_load_kv) {
+      pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytesKV;
+      return MainloopPipeline(shared_storage.pipeline_kv, pipeline_params,
+                              /*cluster_shape=*/Shape<_1, _1, _1>{});
+    } else {
+      return MainloopPipeline(shared_storage.pipeline_kv, pipeline_params);
+    }
+  }();
+  __syncthreads();
+
+  CollectiveMainloop collective_mainloop;
+  CollectiveEpilogue collective_epilogue;
+  
+  if (warp_group_idx == 0) {
+    // producer
+    if constexpr (USE_REG_EALLOC) {
+        cutlass::arch::warpgroup_reg_dealloc<64>();
+    }
+    int const warp_idx = cutlass::canonical_warp_idx_sync();
+    const uint32_t warp_idx_in_warpgroup = __shfl_sync(0xffffffff, warp_idx % 4, 0);
+    PipelineStateQ smem_pipe_write_q = cutlass::make_producer_start_state<MainloopPipelineQ>();
+    PipelineState smem_pipe_write_kv = cutlass::make_producer_start_state<MainloopPipeline>();
+    const int block_id = blockIdx.x;
+    const int bid = mainloop_params.batch_ids[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch ? mainloop_params.tile_ids_per_batch[block_id] : 0;
+    const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
+    const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+    const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+    const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+    // load Q
+    collective_mainloop.load_q(
+        mainloop_params,
+        pipeline_q,
+        smem_pipe_write_q,
+        shared_storage,
+        threadIdx.x,
+        bid,
+        q_tile_id);
+
+    if constexpr (!use_tma_load_kv) {
+      // load kv
+      collective_mainloop.load_kv(
+          mainloop_params,
+          pipeline_kv,
+          smem_pipe_write_kv,
+          shared_storage,
+          bid,
+          seq_len_decoder_now,
+          tile_id
+      );
+    } else {
+      if (warp_idx_in_warpgroup == 0) {
+        // load kv tma
+        collective_mainloop.load_kv_tma(
+            mainloop_params,
+            pipeline_kv,
+            smem_pipe_write_kv,
+            shared_storage,
+            bid,
+            seq_len_decoder_now,
+            tile_id
+        );
+      }
+    }
+  } else if (warp_group_idx == 1) {
+    // mma qk
+    if constexpr (USE_REG_EALLOC) {
+        cutlass::arch::warpgroup_reg_dealloc<64>(); 
+    }
+
+    auto attention_updater = OnlineSoftmax<2 * BLOCK_SHAPE_Q / 64, /*WITH_SCALE=*/true>(mainloop_params.sm_scale);
+    PipelineStateQ smem_pipe_read_q;
+    PipelineState smem_pipe_read_kv;
+    PipelineStateQK smem_pipe_write_qk = cutlass::make_producer_start_state<MainloopPipelineQK>();
+    const int block_id = blockIdx.x;
+    clear(attention_updater.scores_scale);
+    const int bid = mainloop_params.batch_ids[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch ? mainloop_params.tile_ids_per_batch[block_id] : 0;
+    const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
+    const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+    const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+    if constexpr (USE_QK_TWO_STAGE) {
+      mma_qk_two_stages<Ktraits, CAUSAL>(
+          mainloop_params, 
+          pipeline_q, 
+          smem_pipe_read_q,
+          pipeline_qk,
+          smem_pipe_write_qk,
+          pipeline_kv, 
+          smem_pipe_read_kv,
+          attention_updater, 
+          threadIdx.x - NUM_COPY_THREADS,
+          bid,
+          seq_len_decoder_now,
+          seq_len_now,
+          tile_id,
+          q_tile_id,
+          shared_storage);
+    } else {
+      mma_qk_one_stages<Ktraits, CAUSAL>(
+          mainloop_params, 
+          pipeline_q, 
+          smem_pipe_read_q,
+          pipeline_qk,
+          smem_pipe_write_qk,
+          pipeline_kv, 
+          smem_pipe_read_kv,
+          attention_updater, 
+          threadIdx.x - NUM_COPY_THREADS,
+          bid,
+          seq_len_decoder_now,
+          seq_len_now,
+          tile_id,
+          q_tile_id,
+          shared_storage);
+    }
+  } else {
+    // mm pv and store 160 reg
+    if constexpr (USE_REG_EALLOC) {
+        cutlass::arch::warpgroup_reg_alloc<184>();
+    }
+    typename Ktraits::TiledMmaPVSS tiled_mma_pv;
+    Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 1>(TileShape_PDV{}));
+    auto attention_updater = OnlineSoftmax<2 * BLOCK_SHAPE_Q / 64, /*WITH_SCALE=*/true>(mainloop_params.sm_scale);
+    PipelineState smem_pipe_read_kv;
+    PipelineStateQK smem_pipe_read_qk;
+    const int block_id = blockIdx.x;
+    const int bid = mainloop_params.batch_ids[block_id];
+    const int tile_id = mainloop_params.tile_ids_per_batch[block_id];
+    const int q_tile_id = mainloop_params.q_tile_ids_per_batch ? mainloop_params.q_tile_ids_per_batch[block_id] : 0;
+    const int seq_len_now = mainloop_params.seq_lens_this_time[bid];
+    const int seq_len_decoder_now = mainloop_params.seq_lens_decoder[bid] + seq_len_now;
+    const int start_token_idx = mainloop_params.cumsum_q_seqlens[bid];
+    clear(tOrO);
+    clear(attention_updater.scores_scale);
+    mma_pv_one_stages<Ktraits, CAUSAL>(
+        mainloop_params, 
+        pipeline_qk, 
+        smem_pipe_read_qk,
+        pipeline_kv, 
+        smem_pipe_read_kv,
+        attention_updater,
+        tOrO,
+        threadIdx.x - NUM_COPY_THREADS - NUM_MMA_THREADS_QK,
+        bid,
+        seq_len_decoder_now,
+        seq_len_now,
+        tile_id,
+        shared_storage);
+    collective_epilogue.store(
+      epilogue_params, 
+      tOrO, 
+      attention_updater.get_lse(),
+      shared_storage,
+      tiled_mma_pv, 
+      threadIdx.x - NUM_COPY_THREADS - NUM_MMA_THREADS_QK,
+      bid,
+      mainloop_params.bsz,
+      seq_len_now,
+      start_token_idx,
+      tile_id,
+      q_tile_id,
+      seq_len_decoder_now,
+      mainloop_params.chunk_size,
+      mainloop_params.draft_total_token_num,
+      mainloop_params.o_stride_head_num);
+  }
+}
+
+template <typename KernelTraits, bool CAUSAL, typename Params, bool USE_REG_EALLOC=false, bool USE_QK_TWO_STAGE=false>
+cudaError_t BatchMLAWithPagedKVCacheWG4KernelTraitsDispatched(Params& params,
+                                                              cudaStream_t stream) {
+  using DTypeQ = typename KernelTraits::DTypeQ;
+  using DTypeKV = typename KernelTraits::DTypeKV;
+  using DTypeO = typename KernelTraits::DTypeO;
+  using IdType = typename KernelTraits::IdType;
+  using NV_TYPE = typename KernelTraits::NV_TYPE;
+
+  using CollectiveMainloop = CollectiveMainloop<KernelTraits, CAUSAL>;
+  using CollectiveEpilogue = CollectiveEpilogue<KernelTraits>;
+
+  typename CollectiveMainloop::Params mainloop_params = CollectiveMainloop::to_underlying_arguments({
+      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.qk_head_dim), make_stride(params.qk_head_dim, _1{})), // layout q
+      make_layout(make_shape(params.block_size, params.qk_head_dim, params.max_block_num), make_stride(params.qk_head_dim, _1{}, params.block_size * params.qk_head_dim)),
+      make_layout(make_shape(params.chunk_num, params.bsz * params.draft_total_token_num * params.q_num_head), make_stride(params.bsz * params.draft_total_token_num * params.q_num_head, _1{})),
+      params.Q,
+      params.KV,
+      params.m,
+      params.d,
+      params.block_tables,
+      params.seq_lens_this_time,
+      params.seq_lens_encoder,
+      params.seq_lens_decoder,
+      params.cumsum_q_seqlens,
+      params.batch_ids,
+      params.tile_ids_per_batch,
+      params.q_tile_ids_per_batch,
+      params.num_blocks_x,
+      params.sm_scale,
+      params.bsz,
+      params.max_block_num,
+      params.max_block_num_per_seq,
+      params.q_stride_bsz,
+      params.q_stride_head_num,
+      params.kv_stride_block_num,
+      params.kv_stride_block_size,
+      params.o_stride_bsz,
+      params.o_stride_head_num,
+      params.chunk_size,
+      params.chunk_num,
+      params.draft_total_token_num
+  });
+  typename CollectiveEpilogue::Params epilogue_params = CollectiveEpilogue::to_underlying_arguments_ntma({
+      params.O,
+      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.vo_head_dim), make_stride(params.vo_head_dim, _1{})), // layout O
+      params.O_tmp,
+      make_layout(make_shape(KernelTraits::BLOCK_SHAPE_Q, params.vo_head_dim), make_stride(params.vo_head_dim, _1{})) // layout O_tmp
+  });
+
+  // Get the ptr to kernel function.
+  auto kernel =
+      MLAWithKVCacheWG4Kernel<CollectiveMainloop, CollectiveEpilogue, KernelTraits, CAUSAL, USE_REG_EALLOC, USE_QK_TWO_STAGE>;
+  int smem_size = sizeof(typename KernelTraits::SharedStorage);
+  cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  int device;
+  cudaGetDevice(&device);
+  int multiprocessor_count;
+  cudaDeviceGetAttribute(&multiprocessor_count, cudaDevAttrMultiProcessorCount, device);
+  int act_blocks_per_sm;
+  cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+      &act_blocks_per_sm, kernel, KernelTraits::NUM_WARPS * 32, smem_size);
+  dim3 grid_dims = {params.num_blocks_x_int, 1, 1};
+  static constexpr int ctaSize = KernelTraits::NUM_WARPS * 32;
+  dim3 block_dims(ctaSize, 1, 1);
+  kernel<<<grid_dims, block_dims, smem_size, stream>>>(
+    mainloop_params, epilogue_params
+  );
+  if (params.chunk_num > 1) {
+    constexpr int vec_size = 16 / sizeof(DTypeO);
+    constexpr int merge_block_size = 256;
+    constexpr int blockx = KernelTraits::HEAD_DIM_VO / vec_size;
+    constexpr int blocky = (merge_block_size + blockx - 1) / blockx;
+    dim3 grids_merge(min(multiprocessor_count, params.token_num), params.q_num_head); // 128k is too large
+    dim3 blocks_merge(blockx, blocky);
+    merge_multi_chunks_kernel<NV_TYPE, vec_size, blocky, KernelTraits::HEAD_DIM_VO><<<grids_merge, blocks_merge, 0, stream>>>(
+      reinterpret_cast<NV_TYPE*>(params.O_tmp),
+      params.m,
+      params.d,
+      params.seq_lens_this_time,
+      params.seq_lens_decoder,
+      params.seq_lens_encoder,
+      params.padding_offsets,
+      reinterpret_cast<NV_TYPE*>(params.O),
+      params.max_seq_len,
+      params.chunk_num,
+      params.q_num_head,
+      params.chunk_size,
+      params.vo_head_dim,
+      params.token_num,
+      params.bsz,
+      params.draft_total_token_num
+    );
+  }
+  return cudaSuccess;
+}
+
+template <uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO, typename NV_TYPE, typename Params, bool USE_REG_EALLOC=false, bool USE_QK_TWO_STAGE=false>
+cudaError_t BatchMLAWithPagedKVCacheWG4Dispatched(Params& params, cudaStream_t stream) {
+  constexpr bool CAUSAL = true;
+  if constexpr (HEAD_DIM_QK == 576) {
+    if (params.block_size == 32) {
+      DISPATCH_GROUP_SIZE(params.q_num_head, GROUP_SIZE,
+        BatchMLAWithPagedKVCacheWG4KernelTraitsDispatched<
+            AttentionKernelWG4Traits</*USE_TMA_LOAD_KV=*/false, 
+                                     HEAD_DIM_QK, 
+                                     HEAD_DIM_VO, 
+                                     GROUP_SIZE,
+                                     /*BLOCK_SHAPE_Q_=*/64,
+                                     /*BLOCK_SHAPE_KV_=*/32,
+                                     /*NUM_STAGES_=*/4, 
+                                     typename Params::DTypeQ,
+                                     typename Params::DTypeKV, 
+                                     typename Params::DTypeO,
+                                     typename Params::IdType,
+                                     NV_TYPE>,
+            CAUSAL,
+            Params,
+            USE_REG_EALLOC,
+            USE_QK_TWO_STAGE>(params, stream);)
+    }
+  } else {
+    return cudaErrorNotSupported;
+  }
+  return cudaSuccess;
+};
+
+}  // namespace mla_attn
+
+#endif  // ATTENTION_HOPPER_PREFILL_WG4_SM90_CUH_

csrc/gpu/mla_attn/named_barrier.cuh

@@ -0,0 +1,50 @@
+// 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.
+
+/*
+ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri
+ * Dao. Licensed under the BSD 3-Clause.
+ *
+ * Modified by the FlashInfer team.
+ */
+
+#pragma once
+
+#ifndef ATTENTION_HOPPER_NAMED_BARRIERS_CUH_
+#define ATTENTION_HOPPER_NAMED_BARRIERS_CUH_
+
+#include <cuda_runtime.h>
+
+#include "cutlass/arch/barrier.h"
+#include "cutlass/cutlass.h"
+
+namespace mla_attn {
+
+enum class NamedBarriers {
+  kQueryEmpty = 0,
+  kValueEmpty = 1,
+  kWarpSchedulerWG1 = 2,
+  kWarpSchedulerWG2 = 3,
+  kWarpSchedulerWG3 = 4,
+  kPrefetchIndices = 5,
+  kOdone = 6,
+  kWG1WG2Sync = 7,
+  kWG0WG1WG2Sync = 8,
+  kWG1WG2LastSync = 9,
+  kAllWGSync = 10,
+};
+
+}  // namespace mla_attn
+
+#endif  // ATTENTION_HOPPER_NAMED_BARRIERS_CUH_

csrc/gpu/mla_attn/utils.cuh

@@ -0,0 +1,353 @@
+// 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
+
+#ifndef ATTENTION_HOPPER_UTILS_CUH_
+#define ATTENTION_HOPPER_UTILS_CUH_
+
+#include "cute/tensor.hpp"
+#include "cutlass/cutlass.h"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/default_epilogue.hpp"
+#include "cutlass/epilogue/thread/linear_combination.h"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/device/gemm_grouped.h"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/group_array_problem_shape.hpp"
+#include "cutlass/gemm/kernel/default_gemm_grouped.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/numeric_types.h"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/util/command_line.h"
+#include "cutlass/util/distribution.h"
+#include "cutlass/util/host_tensor.h"
+#include "cutlass/util/packed_stride.hpp"
+#include "cutlass/util/reference/device/gemm.h"
+#include "cutlass/util/reference/device/tensor_compare.h"
+#include "cutlass/util/reference/device/tensor_fill.h"
+#include "cutlass/util/tensor_view_io.h"
+#include <assert.h>
+#include <cuda_fp16.h>
+#include <stdint.h>
+#include <stdlib.h>
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+#include <cuda_bf16.h>
+#endif
+
+#include <cuda_runtime.h>
+#include <cutlass/array.h>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/numeric_types.h>
+
+#include <cmath>
+#include <cute/arch/cluster_sm90.hpp>
+#include <cute/tensor.hpp>
+#include "cutlass/fast_math.h"
+
+namespace mla_attn {
+
+using namespace cute;
+
+template <typename TensorT>
+CUTLASS_HOST_DEVICE auto flatten_1(TensorT tensor) {
+  Tensor tensor_flatten = cute::flatten(tensor);
+  return cute::group_modes<1, rank(tensor_flatten)>(tensor_flatten);
+}
+
+CUTLASS_HOST_DEVICE auto get_gmem_layout(int nnz, int num_heads, int head_dim, int64_t n_stride,
+                                         int64_t h_stride) {
+  return make_layout(make_shape(nnz, head_dim, num_heads),
+                     make_stride(n_stride, cute::_1{}, h_stride));
+}
+
+CUTLASS_HOST_DEVICE auto get_lse_gmem_layout(int nnz, int num_heads) {
+  return make_layout(make_shape(num_heads, nnz), make_stride(cute::_1{}, int64_t(num_heads)));
+}
+
+template <typename MTensor, typename Shape>
+CUTLASS_DEVICE auto get_local_tile_tensor(const MTensor& m_tensor, const Shape& tile_shape,
+                                          int head_idx, int offset, int seq_len) {
+  auto g_offset = local_tile(m_tensor(_, _, head_idx), cute::make_shape(1, get<1>(tile_shape)),
+                             make_coord(offset, _0{}));
+  auto g_sequence =
+      make_tensor(g_offset.data(),
+                  make_layout(cute::make_shape(seq_len, get<1>(tile_shape)), g_offset.stride()));
+  auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_, _0{}));
+  return g_tensor;
+}
+
+template <typename MTensor, typename Shape>
+CUTLASS_DEVICE auto get_lse_local_tile_tensor(const MTensor& m_tensor, const Shape& tile_shape,
+                                              int head_idx, int offset, int seq_len) {
+  auto g_offset = local_tile(m_tensor(head_idx, _), cute::make_shape(_1{}), make_coord(offset));
+
+  auto g_sequence = make_tensor(g_offset.data(), make_layout(cute::make_shape(seq_len),
+                                                             cute::make_shape(shape<0>(m_tensor))));
+  auto g_tensor = local_tile(g_sequence, tile_shape, make_coord(_));
+  return g_tensor;
+}
+
+// For SM90, convert acc_layout from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V,
+// MMA_N))
+template <typename Layout>
+__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
+  static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+  static_assert(decltype(size<0, 1>(acc_layout))::value == 2);
+  static_assert(decltype(rank(acc_layout))::value == 3);
+  auto l = acc_layout;
+  return make_layout(make_layout(get<0, 1>(l), get<1>(l)),
+                     make_layout(get<0, 0>(l), get<0, 2>(l), get<2>(l)));
+};
+
+// For SM90, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2, 2), MMA_M, (N / 16,
+// MMA_N))
+template <typename MMA_traits, typename Layout>
+__forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
+  using X = Underscore;
+  static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+  static_assert(decltype(size<0, 1>(acc_layout))::value == 2);
+  static_assert(decltype(rank(acc_layout))::value == 3);
+  static_assert(decltype(rank(get<0>(acc_layout)))::value == 3);
+  auto l = logical_divide(get<0>(acc_layout), Shape<X, X, _2>{});  // (2, 2, (2, N / 16)))
+  return make_layout(make_layout(get<0>(l), get<1>(l), get<2, 0>(l)), get<1>(acc_layout),
+                     make_layout(get<2, 1>(l), get<2>(acc_layout)));
+};
+
+template <typename To_type, typename Engine, typename Layout>
+__forceinline__ __device__ auto convert_type(Tensor<Engine, Layout> const& tensor) {
+  using From_type = typename Engine::value_type;
+  constexpr int numel = decltype(size(tensor))::value;
+  cutlass::NumericArrayConverter<To_type, From_type, numel> convert_op;
+  auto frag = convert_op(*reinterpret_cast<const cutlass::Array<From_type, numel>*>(tensor.data()));
+  return make_tensor(make_rmem_ptr<To_type>(&frag), tensor.layout());
+}
+
+template <bool init = false, int wg_wait = 0, typename TensorA, typename TensorB, typename TensorC,
+          typename TiledMma>
+__forceinline__ __device__ void gemm(TiledMma& tiled_mma, TensorA const& tCrA, TensorB const& tCrB,
+                                     TensorC& tCrC) {
+  constexpr bool Is_RS =
+      !cute::is_base_of<cute::GMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value;
+  // Need to cast away const on tCrA since warpgroup_fence_operand doesn't take const
+  if constexpr (Is_RS) {
+    warpgroup_fence_operand(const_cast<TensorA&>(tCrA));
+  }
+  warpgroup_fence_operand(tCrC);
+  warpgroup_arrive();
+  if constexpr (init) {
+    tiled_mma.accumulate_ = GMMA::ScaleOut::Zero;
+    // Unroll the K mode manually to set scale D to 1
+    CUTLASS_PRAGMA_UNROLL
+    for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+      cute::gemm(tiled_mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC);
+      tiled_mma.accumulate_ = GMMA::ScaleOut::One;
+    }
+  } else {
+    // cute::gemm(tiled_mma, tCrA, tCrB, tCrC);
+    // Unroll the K mode manually to set scale D to 1
+    CUTLASS_PRAGMA_UNROLL
+    for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+      cute::gemm(tiled_mma, tCrA(_, _, k_block), tCrB(_, _, k_block), tCrC);
+      tiled_mma.accumulate_ = GMMA::ScaleOut::One;
+    }
+  }
+  warpgroup_commit_batch();
+  if constexpr (wg_wait >= 0) {
+    warpgroup_wait<wg_wait>();
+  }
+  warpgroup_fence_operand(tCrC);
+  if constexpr (Is_RS) {
+    warpgroup_fence_operand(const_cast<TensorA&>(tCrA));
+  }
+}
+
+#define HOSTDEVICE __host__ __device__
+
+template <typename T, int Size>
+struct alignas(sizeof(T) * Size) AlignedVector {
+  T val[Size];
+
+  HOSTDEVICE inline const T& operator[](int i) const { return val[i]; }
+  HOSTDEVICE inline T& operator[](int i) { return val[i]; }
+};
+
+template <typename T, int Size>
+HOSTDEVICE inline void Load(const T* addr, AlignedVector<T, Size>* vec) {
+  const AlignedVector<T, Size>* addr_vec =
+      reinterpret_cast<const AlignedVector<T, Size>*>(addr);
+  *vec = *addr_vec;
+}
+
+template <typename T, int Size>
+HOSTDEVICE inline void Store(const AlignedVector<T, Size>& vec, T* addr) {
+  AlignedVector<T, Size>* addr_vec =
+      reinterpret_cast<AlignedVector<T, Size>*>(addr);
+  *addr_vec = vec;
+}
+
+template <size_t vec_size, typename T>
+struct prefill_softmax_state_t {
+  AlignedVector<T, vec_size> o;
+  float m;
+  float d;
+  
+  __device__ __forceinline__ void init() {
+    if constexpr (std::is_same<T, half>::value) {
+#pragma unroll
+      for (int i = 0; i < vec_size / 2; ++i) {
+        *((half2*)(&o) + i) = make_half2(0, 0);
+      }
+    } else if constexpr (std::is_same<T, __nv_bfloat16>::value) {
+#pragma unroll
+      for (int i = 0; i < vec_size / 2; ++i) {
+        *((nv_bfloat162*)(&o) + i) = make_bfloat162(0, 0);
+      }
+    }
+    d = 1.f;
+    if constexpr (std::is_same<T, half>::value) {
+      m = -5e4f;
+    } else if constexpr (std::is_same<T, nv_bfloat16>::value) {
+      m = -3.38953e38f;
+    }
+  }
+
+  __device__ __forceinline__ void merge(const AlignedVector<T, vec_size>& other_o, 
+                                        const float other_m,
+                                        const float other_d) {
+    float m_prev = m, d_prev = d;
+    m = max(m_prev, other_m);
+    const float scale1 = __expf(m_prev - m), scale2 = __expf(other_m - m);
+    const T scale1_T = static_cast<T>(scale1), scale2_T = static_cast<T>(scale2);
+    d = d_prev * scale1 + other_d * scale2;
+#pragma unroll
+    for (size_t i = 0; i < vec_size; ++i) {
+      o[i] = o[i] * scale1_T + other_o[i] * scale2_T;
+    }
+  }
+
+  __device__ __forceinline__ void normalize() {
+    const T d_t = static_cast<T>(d);
+#pragma unroll
+    for (size_t i = 0; i < vec_size; ++i) {
+      o[i] /= d_t;
+    }
+  }
+};
+
+template <typename T, int vec_size, uint32_t bdy, uint32_t HEAD_DIM>
+__global__ void merge_multi_chunks_kernel(const T * __restrict__ multi_out, // [num_chunks, bsz, max_draft_token, num_heads, head_dim]
+                                          const float * __restrict__ multi_m, // [num_chunks, bsz, max_draft_token, num_heads]
+                                          const float * __restrict__ multi_d, // [num_chunks, bsz, max_draft_token, num_heads]
+                                          const int * __restrict__ seq_lens_this_time,
+                                          const int * __restrict__ seq_lens_decoder,
+                                          const int * __restrict__ seq_lens_encoder,
+                                          const int * __restrict__ padding_offsets,
+                                          T * __restrict__ out, // [token_num, num_heads, head_dim]
+                                          const int max_seq_len,
+                                          const int num_chunks,
+                                          const int num_heads,
+                                          const int chunk_size,
+                                          const int head_dim,
+                                          const int token_num,
+                                          const int bsz,
+                                          const int draft_total_token_num) {
+  const int vid = threadIdx.x, ty = threadIdx.y;
+  const int hid = blockIdx.y;
+  __shared__ T smem[bdy * HEAD_DIM];
+  __shared__ float md_smem[bdy * 2];
+  for (int qid = blockIdx.x; qid < token_num; qid += gridDim.x) {
+    const uint32_t ori_token_id = qid + padding_offsets[qid];
+    const uint32_t bid = ori_token_id / max_seq_len;
+    const int seq_len_q = seq_lens_this_time[bid];
+    if (seq_len_q == 0) continue;
+    const uint32_t local_seq_id = ori_token_id % max_seq_len;
+    int seq_len_kv = seq_lens_decoder[bid];
+    if (seq_len_kv == 0) continue;
+    seq_len_kv += seq_len_q;
+    const int num_chunks_this_seq = cute::ceil_div(seq_len_kv, chunk_size);
+    if (num_chunks_this_seq <= 1) {
+      // not need merge
+      continue;
+    }
+
+    using LoadT = AlignedVector<T, vec_size>;
+    LoadT load_vec;
+    LoadT res_vec;
+    if constexpr (std::is_same<T, half>::value) {
+#pragma unroll
+      for (int i = 0; i < vec_size / 2; ++i) {
+        *((half2*)(&res_vec) + i) = make_half2(0, 0);
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < vec_size / 2; ++i) {
+        *((nv_bfloat162*)(&res_vec) + i) = make_bfloat162(0, 0);
+      }
+    }
+    float m;
+    float d = 1.f;
+    if constexpr (std::is_same<T, half>::value) {
+      m = -5e4f;
+    } else if constexpr (std::is_same<T, __nv_bfloat16>::value) {
+      m = -3.0e+30f;
+    }
+
+    for (int i = ty; i < num_chunks_this_seq; i += bdy) {
+      uint32_t offset;
+      offset = ((i * bsz + bid) * draft_total_token_num + local_seq_id) * num_heads + hid;
+      float m_prev = m;
+      float d_prev = d;
+      const float m_now = multi_m[offset];
+      const float d_now = multi_d[offset];
+      m = max(m_prev, m_now);
+      offset = (((i * bsz + bid) * draft_total_token_num + local_seq_id) * num_heads + hid) * head_dim + vid * vec_size;
+      Load<T, vec_size>(&multi_out[offset], &load_vec);
+      const float scale1 = __expf(m_prev - m), scale2 = __expf(m_now - m);
+      const T scale1_T = static_cast<T>(scale1), scale2_T = static_cast<T>(scale2);
+      d = d * scale1 + d_now * scale2;
+#pragma unroll
+      for (int j = 0; j < vec_size; j++) {
+        res_vec[j] = res_vec[j] * scale1_T + load_vec[j] * scale2_T;
+      }
+    }
+    // store ty res
+    Store<T, vec_size>(res_vec, &smem[ty * head_dim + vid * vec_size]);
+    md_smem[2 * ty] = m;
+    md_smem[2 * ty + 1] = d;
+    __syncthreads();
+    if (ty == 0) {
+      // merge bdy
+      prefill_softmax_state_t<vec_size, T> st;
+      st.init();
+#pragma unroll
+      for (int i = 0; i < bdy; i++) {
+        Load<T, vec_size>(&smem[i * head_dim + vid * vec_size], &load_vec);
+        const float m_tmp = md_smem[2 * i], d_tmp = md_smem[2 * i + 1];
+        st.merge(load_vec, m_tmp, d_tmp);
+      }
+      st.normalize();
+      Store<T, vec_size>(st.o, &out[(qid * num_heads + hid) * head_dim + vid * vec_size]);
+    }
+    __syncthreads();
+  }
+}
+
+}  // namespace mla_attn
+
+#endif  // ATTENTION_HOPPER_UTILS_CUH_

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_cutlass_kernel.h

@@ -0,0 +1,611 @@
+/***************************************************************************************************
+ * Copyright (c) 2017-2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ *modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright notice,
+ *this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *notice, this list of conditions and the following disclaimer in the
+ *documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the names of its
+ *contributors may be used to endorse or promote products derived from this
+ *software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ *AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ *IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ *DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY DIRECT,
+ *INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+ * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ *DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+ *OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
+ *NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
+ *EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief
+*/
+
+#pragma once
+
+#include "cutlass/complex.h"
+#include "cutlass/cutlass.h"
+#include "cutlass/fast_math.h"
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/matrix_coord.h"
+#include "cutlass/semaphore.h"
+
+#include "cutlass/gemm/kernel/gemm_transpose_operands.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/trace.h"
+
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/gemm/kernel/gemm_moe_problem_visitor.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/tile_interleaved_layout.h"
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// This section exists to that we can use the same kernel code for regular gemm
+// and dequantizing gemms. It will dispatch to the dequantizing gemm if the Mma
+// type has an Iterator for scales in global.
+template <typename...>
+using void_t = void;
+
+template <typename Mma, typename = void>
+struct use_dq_gemm : platform::false_type {};
+
+template <typename Mma>
+struct use_dq_gemm<Mma, void_t<typename Mma::IteratorScale>>
+    : platform::true_type {};
+
+// SFINAE overload for dequantizing gemm
+template <
+    typename Mma,
+    typename ElementScale,
+    typename platform::enable_if<use_dq_gemm<Mma>::value, bool>::type = true>
+CUTLASS_DEVICE static void run_mma(Mma mma,
+                                   int gemm_k_iterations,
+                                   typename Mma::FragmentC& accum,  // NOLINT
+                                   typename Mma::IteratorA iterator_A,
+                                   typename Mma::IteratorB iterator_B,
+                                   typename Mma::FragmentC const& src_accum,
+                                   ElementScale* weight_scale_ptr,
+                                   MatrixCoord scale_extent,
+                                   const int thread_idx,
+                                   MatrixCoord tb_offset_scale) {
+  typename Mma::IteratorScale iterator_scale(
+      Mma::IteratorScale::Layout(scale_extent.column()),
+      weight_scale_ptr,
+      scale_extent,
+      thread_idx,
+      tb_offset_scale);
+
+  mma(gemm_k_iterations,
+      accum,
+      iterator_A,
+      iterator_B,
+      iterator_scale,
+      src_accum);
+}
+
+// SFINAE overload for normal gemm. This completely ignores the scale parameters
+template <
+    typename Mma,
+    typename ElementScale,
+    typename platform::enable_if<!use_dq_gemm<Mma>::value, bool>::type = true>
+CUTLASS_DEVICE static void run_mma(Mma mma,
+                                   int gemm_k_iterations,
+                                   typename Mma::FragmentC& accum,  // NOLINT
+                                   typename Mma::IteratorA iterator_A,
+                                   typename Mma::IteratorB iterator_B,
+                                   typename Mma::FragmentC const& src_accum,
+                                   ElementScale* weight_scale_ptr,
+                                   MatrixCoord scale_extent,
+                                   const int thread_idx,
+                                   MatrixCoord tb_offset_scale) {
+  mma(gemm_k_iterations, accum, iterator_A, iterator_B, src_accum);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Mma_,  ///! Threadblock-scoped matrix multiply-accumulate
+          typename Epilogue_,            ///! Epilogue
+          typename ThreadblockSwizzle_,  ///! Threadblock swizzling function
+          typename KernelArch,  ///! The Architecture this kernel is compiled
+                                /// for. Used since SIMT kernels lose top-level
+                                /// arch.
+          GroupScheduleMode GroupScheduleMode_  ///! Type of scheduling to //
+                                                /// NOLINT perform
+          >
+struct MoeFCGemm {
+ public:
+  using Mma = Mma_;
+  using Epilogue = Epilogue_;
+  using EpilogueOutputOp = typename Epilogue::OutputOp;
+  using ThreadblockSwizzle = ThreadblockSwizzle_;
+  static GroupScheduleMode const kGroupScheduleMode = GroupScheduleMode_;
+  static bool const kTransposed = false;
+
+  // Optional transpose
+  using MapArguments =
+      kernel::detail::MapArguments<typename Mma::IteratorA::Element,
+                                   typename Mma::IteratorA::Layout,
+                                   Mma::kTransformA,
+                                   Mma::IteratorA::AccessType::kElements,
+                                   typename Mma::IteratorB::Element,
+                                   typename Mma::IteratorB::Layout,
+                                   Mma::kTransformB,
+                                   Mma::IteratorB::AccessType::kElements,
+                                   typename Mma::LayoutC,
+                                   kTransposed>;
+
+  // Public-facing type definitions related to operand element type, layout, and
+  // complex conjugate operation. Must interact with the 'kTransposed' notion.
+  static_assert(!kTransposed, "Transpose problem not supported");
+  using ElementA = typename MapArguments::ElementA;
+  using LayoutA = typename MapArguments::LayoutA;
+  using ElementB = typename MapArguments::ElementB;
+  using LayoutB = typename MapArguments::LayoutB;
+  using ElementC = typename Epilogue::OutputTileIterator::Element;
+  using LayoutC = typename MapArguments::LayoutC;
+  using ElementScale = ElementC;
+
+  static ComplexTransform const kTransformA = MapArguments::kTransformA;
+  static ComplexTransform const kTransformB = MapArguments::kTransformB;
+
+  // Type definitions about the mainloop.
+  using Operator = typename Mma::Operator;
+  using OperatorClass = typename Mma::Operator::OperatorClass;
+  using ThreadblockShape = typename Mma::Shape;
+  using WarpShape = typename Mma::Operator::Shape;
+  using InstructionShape = typename Mma::Policy::Operator::InstructionShape;
+  using ArchTag = typename Mma::ArchTag;
+
+  static int const kStages = Mma::kStages;
+  static int const kAlignmentA = MapArguments::kAlignmentA;
+  static int const kAlignmentB = MapArguments::kAlignmentB;
+  static int const kAlignmentC =
+      Epilogue::OutputTileIterator::kElementsPerAccess;
+
+  /// Warp count (concept: GemmShape)
+  using WarpCount = typename Mma::WarpCount;
+  static int const kThreadCount = 32 * WarpCount::kCount;
+
+  using ProblemVisitor = GemmMoeProblemVisitor<ThreadblockShape,
+                                               kGroupScheduleMode,
+                                               kThreadCount,
+                                               kThreadCount,
+                                               kTransposed>;
+
+  //
+  // Structures
+  //
+
+  /// Argument structure
+  struct Arguments {
+    //
+    // Data members
+    //
+
+    int problem_count;
+    int threadblock_count;
+
+    typename EpilogueOutputOp::Params output_op;
+
+    ElementA* ptr_A;
+    ElementB* ptr_B;
+    ElementScale* weight_scales;
+    ElementC* ptr_C;
+    ElementC* ptr_D;
+
+    int64_t* total_rows_before_expert;
+    int64_t gemm_n;
+    int64_t gemm_k;
+
+    // Only used by device-level operator
+    GemmCoord* host_problem_sizes;
+
+    //
+    // Methods
+    //
+
+    /// Default ctor
+    CUTLASS_HOST_DEVICE
+    Arguments()
+        : problem_count(0),
+          threadblock_count(0),
+          ptr_A(nullptr),
+          ptr_B(nullptr),
+          weight_scales(nullptr),
+          ptr_C(nullptr),
+          ptr_D(nullptr),
+          total_rows_before_expert(nullptr),
+          gemm_n(0),
+          gemm_k(0),
+          host_problem_sizes(nullptr) {}
+
+    /// Ctor
+    CUTLASS_HOST_DEVICE
+    Arguments(int problem_count,
+              int threadblock_count,
+              typename EpilogueOutputOp::Params output_op,
+              const ElementA* ptr_A,
+              const ElementB* ptr_B,
+              const ElementScale* weight_scales,
+              const ElementC* ptr_C,
+              ElementC* ptr_D,
+              int64_t* total_rows_before_expert,
+              int64_t gemm_n,
+              int64_t gemm_k,
+              GemmCoord* host_problem_sizes = nullptr)
+        : problem_count(problem_count),
+          threadblock_count(threadblock_count),
+          output_op(output_op),
+          ptr_A(const_cast<ElementA*>(ptr_A)),
+          ptr_B(const_cast<ElementB*>(ptr_B)),
+          weight_scales(const_cast<ElementScale*>(weight_scales)),
+          ptr_C(const_cast<ElementC*>(ptr_C)),
+          ptr_D(ptr_D),
+          total_rows_before_expert(total_rows_before_expert),
+          gemm_n(gemm_n),
+          gemm_k(gemm_k),
+          host_problem_sizes(nullptr) {
+      if (platform::is_same<uint8_t, ElementB>::value ||
+          platform::is_same<uint4b_t, ElementB>::value) {
+        assert(weight_scales);
+      }
+    }
+  };
+
+  //
+  // Structure for precomputing values in host memory and passing to kernels
+  //
+
+  /// Parameters structure
+  struct Params {
+    typename ProblemVisitor::Params problem_visitor;
+    int threadblock_count;
+
+    typename EpilogueOutputOp::Params output_op;
+
+    ElementA* ptr_A;
+    ElementB* ptr_B;
+    ElementScale* weight_scales;
+    ElementC* ptr_C;
+    ElementC* ptr_D;
+
+    //
+    // Methods
+    //
+
+    CUTLASS_HOST_DEVICE
+    Params()
+        : ptr_A(nullptr),
+          ptr_B(nullptr),
+          weight_scales(nullptr),
+          ptr_C(nullptr),
+          ptr_D(nullptr) {}
+
+    CUTLASS_HOST_DEVICE
+    Params(Arguments const& args,
+           void* workspace = nullptr,
+           int tile_count = 0)  // NOLINT
+        : problem_visitor(args.total_rows_before_expert,
+                          args.gemm_n,
+                          args.gemm_k,
+                          args.problem_count,
+                          workspace,
+                          tile_count),
+          threadblock_count(args.threadblock_count),
+          output_op(args.output_op),
+          ptr_A(args.ptr_A),
+          ptr_B(args.ptr_B),
+          weight_scales(args.weight_scales),
+          ptr_C(args.ptr_C),
+          ptr_D(args.ptr_D) {}
+
+    CUTLASS_HOST_DEVICE
+    void update(Arguments const& args,
+                void* workspace = nullptr,
+                int tile_count = 0) {
+      problem_visitor =
+          typename ProblemVisitor::Params(args.total_rows_before_expert,
+                                          args.gemm_n,
+                                          args.gemm_k,
+                                          args.problem_count,
+                                          workspace,
+                                          tile_count);
+      threadblock_count = args.threadblock_count;
+      output_op = args.output_op;
+      ptr_A = args.ptr_A;
+      ptr_B = args.ptr_B;
+      weight_scales = args.weight_scales;
+      ptr_C = args.ptr_C;
+      ptr_D = args.ptr_D;
+    }
+  };
+
+  /// Shared memory storage structure
+  union SharedStorage {
+    typename ProblemVisitor::SharedStorage problem_visitor;
+    typename Mma::SharedStorage main_loop;
+    typename Epilogue::SharedStorage epilogue;
+  };
+
+ public:
+  //
+  // Methods
+  //
+
+  CUTLASS_DEVICE
+  MoeFCGemm() {}
+
+  /// Determines whether kernel satisfies alignment
+  static Status can_implement(cutlass::gemm::GemmCoord const& problem_size) {
+    return Status::kSuccess;
+  }
+
+  static Status can_implement(Arguments const& args) {
+    if (platform::is_same<uint8_t, ElementB>::value ||
+        platform::is_same<uint4b_t, ElementB>::value) {
+      if (args.weight_scales == nullptr) {
+        CUTLASS_TRACE_HOST(
+            "MoeFCGemm::can_implement() - weight scales are required for "
+            "uint8_t and uint4b_t");
+        return Status::kInvalid;
+      }
+    } else if (args.weight_scales != nullptr) {
+      CUTLASS_TRACE_HOST(
+          "MoeFCGemm::can_implement() - weight scales are ignored for all "
+          "types except uint8_t and uint4b_t");
+      return Status::kInvalid;
+    }
+    return Status::kSuccess;
+  }
+
+  static size_t get_extra_workspace_size(
+      Arguments const& args, cutlass::gemm::GemmCoord const& grid_tiled_shape) {
+    return 0;
+  }
+
+  // The dummy template parameter is not used and exists so that we can compile
+  // this code using a standard earlier than C++17. Prior to C++17, fully
+  // specialized templates HAD to exists in a namespace
+  template <bool B, typename dummy = void>
+  struct KernelRunner {
+    CUTLASS_DEVICE
+    static void run_kernel(Params const& params,
+                           SharedStorage& shared_storage) {  // NOLINT
+      CUTLASS_NOT_IMPLEMENTED();
+    }
+  };
+
+  template <typename dummy>
+  struct KernelRunner<true, dummy> {
+    CUTLASS_DEVICE
+    static void run_kernel(Params const& params,
+                           SharedStorage& shared_storage) {  // NOLINT
+      //
+      // These types shadow the type-level definitions and support the ability
+      // to implement a 'transposed' GEMM that computes the transposed problems.
+      //
+      using ElementA = typename Mma::IteratorA::Element;
+      using LayoutA = typename Mma::IteratorA::Layout;
+      using ElementB = typename Mma::IteratorB::Element;
+      using LayoutB = typename Mma::IteratorB::Layout;
+      using ElementC = typename Epilogue::OutputTileIterator::Element;
+      using LayoutC = typename Epilogue::OutputTileIterator::Layout;
+      static constexpr int kInterleave =
+          Mma::IteratorB::Shape::kRow / Mma::Shape::kK;
+      static_assert(
+          platform::is_same<LayoutB, layout::RowMajor>::value &&
+                  kInterleave == 1 ||
+              platform::is_same<LayoutB, layout::ColumnMajor>::value &&
+                  kInterleave >= 1,
+          "B must be row major/col major OR col major interleaved.");
+
+      //
+      // Problem visitor.
+      //
+      ProblemVisitor problem_visitor(
+          params.problem_visitor, shared_storage.problem_visitor, blockIdx.x);
+
+      const int64_t gemm_k = params.problem_visitor.gemm_k;
+      const int64_t gemm_n = params.problem_visitor.gemm_n;
+      int64_t bytes_per_expert_matrix =
+          (gemm_k * gemm_n / 8) * cutlass::sizeof_bits<ElementB>::value;
+
+      // Outer 'persistent' loop to iterate over tiles
+      while (problem_visitor.next_tile()) {
+        GemmCoord problem_size = problem_visitor.problem_size();
+        int32_t problem_idx = problem_visitor.problem_index();
+        int32_t cta_idx = int32_t(problem_visitor.threadblock_idx());
+
+        GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
+
+        cutlass::gemm::GemmCoord threadblock_offset(
+            int(cta_idx / grid_shape.n()) * Mma::Shape::kM,  // NOLINT
+            int(cta_idx % grid_shape.n()) * Mma::Shape::kN,  // NOLINT
+            0);
+
+        // Load element pointers. Exchange pointers and strides if working on
+        // the transpose
+        const int64_t rows_to_jump =
+            problem_idx == 0
+                ? 0
+                : params.problem_visitor.last_row_for_problem[problem_idx - 1];
+        ElementA* ptr_A =
+            reinterpret_cast<ElementA*>(params.ptr_A) + rows_to_jump * gemm_k;
+        typename LayoutA::LongIndex ldm_A = gemm_k;
+
+        char* byte_ptr_B = ((char*)params.ptr_B) +                 // NOLINT
+                           problem_idx * bytes_per_expert_matrix;  // NOLINT
+        ElementB* ptr_B = reinterpret_cast<ElementB*>(byte_ptr_B);
+        typename LayoutB::LongIndex ldm_B =
+            platform::is_same<layout::RowMajor, LayoutB>::value
+                ? gemm_n
+                : gemm_k * kInterleave;
+
+        // Compute initial location in logical coordinates
+        cutlass::MatrixCoord tb_offset_A{
+            threadblock_offset.m(),
+            0,
+        };
+
+        cutlass::MatrixCoord tb_offset_B{0,
+                                         threadblock_offset.n() / kInterleave};
+
+        cutlass::MatrixCoord tb_offset_scale{0, threadblock_offset.n()};
+
+        // Compute position within threadblock
+        int thread_idx = threadIdx.x;
+
+        // Construct iterators to A and B operands
+        typename Mma::IteratorA iterator_A(LayoutA(ldm_A),
+                                           ptr_A,
+                                           {problem_size.m(), problem_size.k()},
+                                           thread_idx,
+                                           tb_offset_A);
+
+        typename Mma::IteratorB iterator_B(
+            LayoutB(ldm_B),
+            ptr_B,
+            {problem_size.k() * kInterleave, problem_size.n() / kInterleave},
+            thread_idx,
+            tb_offset_B);
+
+        typename Mma::FragmentC accumulators;
+
+        accumulators.clear();
+
+        // Broadcast the warp_id computed by lane 0 to ensure dependent code
+        // is compiled as warp-uniform.
+        int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
+
+        int lane_idx = threadIdx.x % 32;
+
+        //
+        // Matrix multiply phase
+        //
+
+        // Construct thread-scoped matrix multiply
+        auto CreateMMA = [&]() {
+          if constexpr (use_dq_gemm<Mma>::value)
+            return Mma(
+                shared_storage.main_loop, -1, thread_idx, warp_idx, lane_idx);
+          else
+            return Mma(
+                shared_storage.main_loop, thread_idx, warp_idx, lane_idx);
+        };
+        Mma mma = CreateMMA();
+
+        // Compute threadblock-scoped matrix multiply-add
+        int gemm_k_iterations =
+            (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK;
+
+        // Wait for all threads to finish their epilogue phases from the
+        // previous tile.
+        __syncthreads();
+
+        // Compute threadblock-scoped matrix multiply-add
+        ElementScale* weight_scale_ptr =
+            params.weight_scales + problem_idx * problem_size.n();
+        run_mma<Mma>(mma,
+                     gemm_k_iterations,
+                     accumulators,
+                     iterator_A,
+                     iterator_B,
+                     accumulators,
+                     weight_scale_ptr,
+                     {1, problem_size.n()},
+                     thread_idx,
+                     tb_offset_scale);
+
+        //
+        // Epilogue
+        //
+
+        EpilogueOutputOp output_op(params.output_op);
+
+        ElementC* ptr_C =
+            reinterpret_cast<ElementC*>(params.ptr_C) + problem_idx * gemm_n;
+        ElementC* ptr_D =
+            reinterpret_cast<ElementC*>(params.ptr_D) + rows_to_jump * gemm_n;
+
+        LayoutC layout_C(0);
+        LayoutC layout_D(gemm_n);
+
+        typename Epilogue::OutputTileIterator::Params params_C(layout_C);
+        typename Epilogue::OutputTileIterator::Params params_D(layout_D);
+
+        // Tile iterator loading from source tensor.
+        typename Epilogue::OutputTileIterator iterator_C(
+            params_C,
+            ptr_C,
+            problem_size.mn(),
+            thread_idx,
+            threadblock_offset.mn());
+
+        // Tile iterator writing to destination tensor.
+        typename Epilogue::OutputTileIterator iterator_D(
+            params_D,
+            ptr_D,
+            problem_size.mn(),
+            thread_idx,
+            threadblock_offset.mn());
+
+        Epilogue epilogue(
+            shared_storage.epilogue, thread_idx, warp_idx, lane_idx);
+
+        // Execute the epilogue operator to update the destination tensor.
+        epilogue(output_op, iterator_D, accumulators, iterator_C);
+
+        // Next tile
+        problem_visitor.advance(gridDim.x);
+      }
+    }
+  };
+
+  /*
+    To improve compilation speed, we do not compile the device operator if the
+    CUDA_ARCH does not correspond to the ArchTag of the cutlass kernel operator.
+  */
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const& params,
+                  SharedStorage& shared_storage) {  // NOLINT
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700) && (__CUDA_ARCH__ < 750)
+    static constexpr bool compile_needed =
+        platform::is_same<KernelArch, arch::Sm70>::value;
+    KernelRunner<compile_needed>::run_kernel(params, shared_storage);
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 750) && (__CUDA_ARCH__ < 800)
+    static constexpr bool compile_needed =
+        platform::is_same<KernelArch, arch::Sm75>::value;
+    KernelRunner<compile_needed>::run_kernel(params, shared_storage);
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) && (__CUDA_ARCH__ < 900)
+    static constexpr bool compile_needed =
+        platform::is_same<KernelArch, arch::Sm80>::value;
+    KernelRunner<compile_needed>::run_kernel(params, shared_storage);
+#elif defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 910)
+    static constexpr bool compile_needed =
+        platform::is_same<KernelArch, arch::Sm80>::value;
+    KernelRunner<compile_needed>::run_kernel(params, shared_storage);
+#else
+    CUTLASS_NOT_IMPLEMENTED();
+#endif
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h

@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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_runtime_api.h>
+#include <string>
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/ft_gemm_configs.h"
+using namespace phi;
+
+template <typename T, /*The type used for activations/scales/compute*/
+          typename WeightType /* The type for the MoE weights */>
+class MoeGemmRunner {
+ public:
+  MoeGemmRunner();
+
+  void moe_gemm_bias_act(const T* A,
+                         const WeightType* B,
+                         const T* weight_scales,
+                         const T* biases,
+                         T* C,
+                         int64_t* total_rows_before_expert,
+                         int64_t total_rows,
+                         int64_t gemm_n,
+                         int64_t gemm_k,
+                         int num_experts,
+                         std::string activation_type,
+                         cudaStream_t stream);
+
+  void moe_gemm(const T* A,
+                const WeightType* B,
+                const T* weight_scales,
+                T* C,
+                int64_t* total_rows_before_expert,
+                int64_t total_rows,
+                int64_t gemm_n,
+                int64_t gemm_k,
+                int num_experts,
+                cudaStream_t stream);
+
+ private:
+  template <typename EpilogueTag>
+  void dispatch_to_arch(const T* A,
+                        const WeightType* B,
+                        const T* weight_scales,
+                        const T* biases,
+                        T* C,
+                        int64_t* total_rows_before_expert,
+                        int64_t total_rows,
+                        int64_t gemm_n,
+                        int64_t gemm_k,
+                        int num_experts,
+                        CutlassGemmConfig gemm_config,
+                        cudaStream_t stream,
+                        int* occupancy = nullptr);
+
+  template <typename EpilogueTag>
+  void run_gemm(const T* A,
+                const WeightType* B,
+                const T* weight_scales,
+                const T* biases,
+                T* C,
+                int64_t* total_rows_before_expert,
+                int64_t total_rows,
+                int64_t gemm_n,
+                int64_t gemm_k,
+                int num_experts,
+                cudaStream_t stream);
+
+ private:
+  int sm_;
+  int multi_processor_count_;
+};

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_bf16.cu

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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 "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+#ifdef PADDLE_CUDA_BF16
+template class MoeGemmRunner<__nv_bfloat16, __nv_bfloat16>;
+#endif

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int4.cu

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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 "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+#ifdef PADDLE_CUDA_BF16
+template class MoeGemmRunner<__nv_bfloat16, cutlass::uint4b_t>;
+#endif
+

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_bf16_int8.cu

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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 "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+#ifdef PADDLE_CUDA_BF16
+template class MoeGemmRunner<__nv_bfloat16, uint8_t>;
+#endif
+

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_fp16.cu

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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 "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+template class MoeGemmRunner<half, half>;

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int4.cu

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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 "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+template class MoeGemmRunner<half, cutlass::uint4b_t>;

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_fp16_int8.cu

@@ -0,0 +1,23 @@
+/*
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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 "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h"
+#include "helper.h"
+
+
+template class MoeGemmRunner<half, uint8_t>;

csrc/gpu/moe/fused_moe/cutlass_kernels/moe_gemm/fused_moe_gemm_kernels_template.h

@@ -0,0 +1,803 @@
+/*
+ * Copyright (c) 2020-2024, NVIDIA CORPORATION.  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.
+ */
+
+// Ignore CUTLASS warnings about type punning
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#pragma once
+
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <math.h>
+#include <optional>
+#include <sstream>
+#include "cutlass/array.h"
+#include "cutlass/numeric_conversion.h"
+
+#include "cutlass/gemm/device/gemm_grouped.h"
+#include "cutlass/gemm/kernel/default_gemm_grouped.h"
+
+#include "paddle/common/errors.h"
+#include "paddle/phi/core/enforce.h"
+
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/compute_occupancy.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/epilogue_helpers.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/gemm/kernel/default_fpA_intB_traits.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_extensions/gemm/threadblock/default_mma.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_cutlass_kernel.h"
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+
+#pragma GCC diagnostic pop
+
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_kernels/cutlass_heuristic.h"
+#include "paddle/phi/kernels/fusion/cutlass/cutlass_kernels/gemm_config_manager.h"
+
+#include "helper.h"
+
+using namespace phi;
+// ============================= Variable batched Gemm things
+// ===========================
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape,
+          int Stages>
+void generic_moe_gemm_kernelLauncher(const T* A,
+                                     const WeightType* B,
+                                     const T* weight_scales,
+                                     const T* biases,
+                                     T* C,
+                                     int64_t* total_rows_before_expert,
+                                     int64_t gemm_n,
+                                     int64_t gemm_k,
+                                     int num_experts,
+                                     CutlassGemmConfig gemm_config,
+                                     const int multi_processor_count,
+                                     cudaStream_t stream,
+                                     int* kernel_occupancy = nullptr) {
+  if (gemm_config.split_k_style != SplitKStyle::NO_SPLIT_K) {
+    PADDLE_FATAL("[MoeGemm] Grouped gemm does not support split-k");
+  }
+
+#ifdef PADDLE_CUDA_BF16
+  static_assert(cutlass::platform::is_same<T, __nv_bfloat16>::value ||
+                    cutlass::platform::is_same<T, half>::value ||
+                    cutlass::platform::is_same<T, float>::value,
+                "Specialized for bfloat16, half, float");
+#else
+  static_assert(cutlass::platform::is_same<T, half>::value ||
+                    cutlass::platform::is_same<T, float>::value,
+                "Specialized for half, float");
+#endif
+
+  static_assert(
+      cutlass::platform::is_same<T, WeightType>::value ||
+          cutlass::platform::is_same<WeightType, uint8_t>::value ||
+          cutlass::platform::is_same<WeightType, cutlass::uint4b_t>::value,
+      "");
+
+  // The cutlass type for the input elements. This is needed to convert to
+  // cutlass::half_t if necessary.
+  using ElementType_ = typename cutlass::platform::conditional<
+      cutlass::platform::is_same<T, half>::value,
+      cutlass::half_t,
+      T>::type;
+#ifdef PADDLE_CUDA_BF16
+  using ElementType = typename cutlass::platform::conditional<
+      cutlass::platform::is_same<ElementType_, __nv_bfloat16>::value,
+      cutlass::bfloat16_t,
+      ElementType_>::type;
+#else
+  using ElementType = ElementType_;
+#endif
+
+  using CutlassWeightType_ = typename cutlass::platform::conditional<
+      cutlass::platform::is_same<WeightType, half>::value,
+      cutlass::half_t,
+      WeightType>::type;
+#ifdef PADDLE_CUDA_BF16
+  using CutlassWeightType = typename cutlass::platform::conditional<
+      cutlass::platform::is_same<CutlassWeightType_, __nv_bfloat16>::value,
+      cutlass::bfloat16_t,
+      CutlassWeightType_>::type;
+#else
+  using CutlassWeightType = CutlassWeightType_;
+#endif
+
+  // We need separate config for each architecture since we will target
+  // different tensorcore instructions. For float, we do not target TCs.
+  using MixedGemmArchTraits = cutlass::gemm::kernel::
+      MixedGemmArchTraits<ElementType, CutlassWeightType, arch>;
+  using ElementAccumulator = typename MixedGemmArchTraits::AccType;
+
+  using EpilogueOp = typename Epilogue<ElementType,
+                                       MixedGemmArchTraits::ElementsPerAccessC,
+                                       ElementAccumulator,
+                                       EpilogueTag>::Op;
+
+  // Finally, set up the kernel.
+  using GemmKernel_ = typename cutlass::gemm::kernel::DefaultGemmGrouped<
+      ElementType,
+      cutlass::layout::RowMajor,
+      cutlass::ComplexTransform::kNone,
+      MixedGemmArchTraits::ElementsPerAccessA,
+      CutlassWeightType,
+      typename MixedGemmArchTraits::LayoutB,
+      cutlass::ComplexTransform::kNone,
+      MixedGemmArchTraits::ElementsPerAccessB,
+      ElementType,
+      cutlass::layout::RowMajor,
+      ElementAccumulator,
+      typename MixedGemmArchTraits::OperatorClass,
+      arch,
+      ThreadblockShape,
+      WarpShape,
+      typename MixedGemmArchTraits::InstructionShape,
+      EpilogueOp,
+      cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
+      Stages,
+      cutlass::gemm::kernel::GroupScheduleMode::kDeviceOnly,
+      typename MixedGemmArchTraits::Operator>::GemmKernel;
+
+  using GemmKernel =
+      cutlass::gemm::kernel::MoeFCGemm<typename GemmKernel_::Mma,
+                                       typename GemmKernel_::Epilogue,
+                                       typename GemmKernel_::ThreadblockSwizzle,
+                                       arch,  // Ensure top level arch is used
+                                              // for dispatch
+                                       GemmKernel_::kGroupScheduleMode>;
+
+  using GemmGrouped = cutlass::gemm::device::GemmGrouped<GemmKernel>;
+
+  if (kernel_occupancy != nullptr) {
+    *kernel_occupancy = compute_occupancy_for_kernel<GemmKernel>();
+    return;
+  }
+  int occupancy = std::min(2, GemmGrouped::maximum_active_blocks());
+
+  if (occupancy == 0) {
+    PADDLE_FATAL(
+        "[MoE Runner] GPU lacks the shared memory resources to run "
+        "GroupedGEMM kernel");
+  }
+  const int threadblock_count = multi_processor_count * occupancy;
+
+  typename EpilogueOp::Params epilogue_op(ElementAccumulator(1.f),
+                                          ElementAccumulator(0.f));
+
+  typename GemmGrouped::Arguments args(
+      num_experts,
+      threadblock_count,
+      epilogue_op,
+      reinterpret_cast<const ElementType*>(A),
+      reinterpret_cast<const CutlassWeightType*>(B),
+      reinterpret_cast<const ElementType*>(weight_scales),
+      reinterpret_cast<const ElementType*>(biases),
+      reinterpret_cast<ElementType*>(C),
+      total_rows_before_expert,
+      gemm_n,
+      gemm_k);
+
+  GemmGrouped gemm;
+
+  auto can_implement = gemm.can_implement(args);
+  if (can_implement != cutlass::Status::kSuccess) {
+    std::string err_msg = "MoEFC kernel will fail for params. Error: " +
+                          std::string(cutlassGetStatusString(can_implement));
+    PADDLE_FATAL("[MoE Runner] " + err_msg);
+  }
+
+  auto init_status = gemm.initialize(args);
+  if (init_status != cutlass::Status::kSuccess) {
+    std::string err_msg =
+        "Failed to initialize cutlass variable batched gemm. Error: " +
+        std::string(cutlassGetStatusString(init_status));
+    PADDLE_FATAL("[MoE Runner] " + err_msg);
+  }
+
+  auto run_status = gemm.run(stream);
+  if (run_status != cutlass::Status::kSuccess) {
+    std::string err_msg =
+        "Failed to run cutlass variable batched gemm. Error: " +
+        std::string(cutlassGetStatusString(run_status));
+    PADDLE_FATAL("[MoE Runner] " + err_msg);
+  }
+}
+
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape,
+          int Stages,
+          typename Enable = void>
+struct dispatch_stages {
+  static void dispatch(const T* A,
+                       const WeightType* B,
+                       const T* weight_scales,
+                       const T* biases,
+                       T* C,
+                       int64_t* total_rows_before_expert,
+                       int64_t gemm_n,
+                       int64_t gemm_k,
+                       int num_experts,
+                       CutlassGemmConfig gemm_config,
+                       int multi_processor_count,
+                       cudaStream_t stream,
+                       int* occupancy = nullptr) {
+    // FT_LOG_DEBUG(__PRETTY_FUNCTION__);
+    std::string err_msg = "Cutlass fpA_intB gemm. Not instantiates for arch " +
+                          std::to_string(arch::kMinComputeCapability) +
+                          " with stages set to " + std::to_string(Stages);
+    PADDLE_FATAL("[dispatch_stages::dispatch] " + err_msg);
+  }
+};
+
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape>
+struct dispatch_stages<T,
+                       WeightType,
+                       arch,
+                       EpilogueTag,
+                       ThreadblockShape,
+                       WarpShape,
+                       2> {
+  static void dispatch(const T* A,
+                       const WeightType* B,
+                       const T* weight_scales,
+                       const T* biases,
+                       T* C,
+                       int64_t* total_rows_before_expert,
+                       int64_t gemm_n,
+                       int64_t gemm_k,
+                       int num_experts,
+                       CutlassGemmConfig gemm_config,
+                       int multi_processor_count,
+                       cudaStream_t stream,
+                       int* occupancy = nullptr) {
+    generic_moe_gemm_kernelLauncher<T,
+                                    WeightType,
+                                    arch,
+                                    EpilogueTag,
+                                    ThreadblockShape,
+                                    WarpShape,
+                                    2>(A,
+                                       B,
+                                       weight_scales,
+                                       biases,
+                                       C,
+                                       total_rows_before_expert,
+                                       gemm_n,
+                                       gemm_k,
+                                       num_experts,
+                                       gemm_config,
+                                       multi_processor_count,
+                                       stream,
+                                       occupancy);
+  }
+};
+
+template <typename T,
+          typename WeightType,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape,
+          int Stages>
+struct dispatch_stages<T,
+                       WeightType,
+                       cutlass::arch::Sm80,
+                       EpilogueTag,
+                       ThreadblockShape,
+                       WarpShape,
+                       Stages,
+                       typename std::enable_if<(Stages > 2)>::type> {
+  static void dispatch(const T* A,
+                       const WeightType* B,
+                       const T* weight_scales,
+                       const T* biases,
+                       T* C,
+                       int64_t* total_rows_before_expert,
+                       int64_t gemm_n,
+                       int64_t gemm_k,
+                       int num_experts,
+                       CutlassGemmConfig gemm_config,
+                       int multi_processor_count,
+                       cudaStream_t stream,
+                       int* occupancy = nullptr) {
+    generic_moe_gemm_kernelLauncher<T,
+                                    WeightType,
+                                    cutlass::arch::Sm80,
+                                    EpilogueTag,
+                                    ThreadblockShape,
+                                    WarpShape,
+                                    Stages>(A,
+                                            B,
+                                            weight_scales,
+                                            biases,
+                                            C,
+                                            total_rows_before_expert,
+                                            gemm_n,
+                                            gemm_k,
+                                            num_experts,
+                                            gemm_config,
+                                            multi_processor_count,
+                                            stream,
+                                            occupancy);
+  }
+};
+
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename ThreadblockShape,
+          typename WarpShape>
+void dispatch_gemm_config(const T* A,
+                          const WeightType* B,
+                          const T* weight_scales,
+                          const T* biases,
+                          T* C,
+                          int64_t* total_rows_before_expert,
+                          int64_t gemm_n,
+                          int64_t gemm_k,
+                          int num_experts,
+                          CutlassGemmConfig gemm_config,
+                          int multi_processor_count,
+                          cudaStream_t stream,
+                          int* occupancy = nullptr) {
+#define dispatch_stages_macro(STAGE)                           \
+  case STAGE:                                                  \
+    dispatch_stages<T,                                         \
+                    WeightType,                                \
+                    arch,                                      \
+                    EpilogueTag,                               \
+                    ThreadblockShape,                          \
+                    WarpShape,                                 \
+                    STAGE>::dispatch(A,                        \
+                                     B,                        \
+                                     weight_scales,            \
+                                     biases,                   \
+                                     C,                        \
+                                     total_rows_before_expert, \
+                                     gemm_n,                   \
+                                     gemm_k,                   \
+                                     num_experts,              \
+                                     gemm_config,              \
+                                     multi_processor_count,    \
+                                     stream,                   \
+                                     occupancy);               \
+    break;
+
+  switch (gemm_config.stages) {
+    dispatch_stages_macro(2);
+    dispatch_stages_macro(3);
+    dispatch_stages_macro(4);
+    dispatch_stages_macro(5);
+    default:
+      std::string err_msg = "dispatch_gemm_config does not support stages " +
+                            std::to_string(gemm_config.stages);
+      PADDLE_FATAL("[MoE][dispatch_gemm_config] " + err_msg);
+      break;
+  }
+}
+
+#define dispatch_gemm_config_macro(AA, BB, CC, DD, EE, FF)      \
+  case CutlassTileConfig::                                      \
+      CtaShape##AA##x##BB##x##CC##_WarpShape##DD##x##EE##x##FF: \
+    dispatch_gemm_config<T,                                     \
+                         WeightType,                            \
+                         arch,                                  \
+                         EpilogueTag,                           \
+                         cutlass::gemm::GemmShape<AA, BB, CC>,  \
+                         cutlass::gemm::GemmShape<DD, EE, FF>>( \
+        A,                                                      \
+        B,                                                      \
+        weight_scales,                                          \
+        biases,                                                 \
+        C,                                                      \
+        total_rows_before_expert,                               \
+        gemm_n,                                                 \
+        gemm_k,                                                 \
+        num_experts,                                            \
+        gemm_config,                                            \
+        multi_processor_count,                                  \
+        stream,                                                 \
+        occupancy);                                             \
+    break;
+
+// This overload will handle tensorop gemms. It is disabled via SFINAE for fp32.
+// This overload is only enabled when T == WeightType.
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename std::enable_if<!std::is_same<T, float>::value &&
+                                  std::is_same<T, WeightType>::value>::type* =
+              nullptr>
+void dispatch_moe_gemm_to_cutlass(const T* A,
+                                  const WeightType* B,
+                                  const T* weight_scales,
+                                  const T* biases,
+                                  T* C,
+                                  int64_t* total_rows_before_expert,
+                                  int64_t total_rows,
+                                  int64_t gemm_n,
+                                  int64_t gemm_k,
+                                  int num_experts,
+                                  CutlassGemmConfig gemm_config,
+                                  int sm_version,
+                                  int multi_processor_count,
+                                  cudaStream_t stream,
+                                  int* occupancy = nullptr) {
+  switch (gemm_config.tile_config) {
+    dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
+    dispatch_gemm_config_macro(64, 128, 64, 32, 64, 64);
+    dispatch_gemm_config_macro(128, 128, 64, 64, 32, 64);
+    case CutlassTileConfig::Undefined:
+      PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+      break;
+    case CutlassTileConfig::ChooseWithHeuristic:
+      PADDLE_FATAL(
+          "[dispatch_moe_gemm_to_cutlass] gemm config should have "
+          "already been set by heuristic.");
+      break;
+    default:
+      PADDLE_FATAL(
+          "[dispatch_moe_gemm_to_cutlass] Config is invalid for same "
+          "type MoE tensorop GEMM.");
+      break;
+  }
+}
+
+// Tensorop GEMM overload
+// Overload for quantize MoE GEMMs. We disable some warp configs here since they
+// will not be used and we can improve compile time
+template <typename T,
+          typename WeightType,
+          typename arch,
+          typename EpilogueTag,
+          typename std::enable_if<!std::is_same<T, float>::value &&
+                                  !std::is_same<T, WeightType>::value>::type* =
+              nullptr>
+void dispatch_moe_gemm_to_cutlass(const T* A,
+                                  const WeightType* B,
+                                  const T* weight_scales,
+                                  const T* biases,
+                                  T* C,
+                                  int64_t* total_rows_before_expert,
+                                  int64_t total_rows,
+                                  int64_t gemm_n,
+                                  int64_t gemm_k,
+                                  int num_experts,
+                                  CutlassGemmConfig gemm_config,
+                                  int sm_version,
+                                  int multi_processor_count,
+                                  cudaStream_t stream,
+                                  int* occupancy = nullptr) {
+  if constexpr (std::is_same<arch, cutlass::arch::Sm70>::value) {
+    switch (gemm_config.tile_config) {
+      dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
+      dispatch_gemm_config_macro(64, 128, 64, 64, 64, 64);
+      case CutlassTileConfig::Undefined:
+        PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+        break;
+      case CutlassTileConfig::ChooseWithHeuristic:
+        PADDLE_FATAL(
+            "[dispatch_moe_gemm_to_cutlass] gemm config should have "
+            "already been set by heuristic.");
+        break;
+      default:
+        PADDLE_FATAL(
+            "[dispatch_moe_gemm_to_cutlass] Config is invalid for "
+            "mixed type tensorop GEMM.");
+        break;
+    }
+  } else {
+    switch (gemm_config.tile_config) {
+      dispatch_gemm_config_macro(16, 128, 64, 16, 32, 64);
+      dispatch_gemm_config_macro(32, 128, 64, 32, 32, 64);
+      dispatch_gemm_config_macro(64, 128, 64, 64, 64, 64);
+      dispatch_gemm_config_macro(128, 128, 64, 64, 64, 64);
+      dispatch_gemm_config_macro(128, 128, 64, 128, 32, 64);
+      dispatch_gemm_config_macro(128, 256, 64, 64, 64, 64);
+      dispatch_gemm_config_macro(64, 128, 64, 64, 32, 64);
+      case CutlassTileConfig::Undefined:
+        PADDLE_FATAL("[dispatch_moe_gemm_to_cutlass] gemm config undefined.");
+        break;
+      case CutlassTileConfig::ChooseWithHeuristic:
+        PADDLE_FATAL(
+            "[dispatch_moe_gemm_to_cutlass] gemm config should have "
+            "already been set by heuristic.");
+        break;
+      default:
+        PADDLE_FATAL(
+            "[dispatch_moe_gemm_to_cutlass] Config is invalid for "
+            "mixed type tensorop GEMM.");
+        break;
+    }
+  }
+}
+
+// This overload will handle simt gemms. It is disabled via SFINAE for tensorop.
+template <
+    typename T,
+    typename WeightType,
+    typename arch,
+    typename EpilogueTag,
+    typename std::enable_if<std::is_same<T, float>::value>::type* = nullptr>
+void dispatch_moe_gemm_to_cutlass(const T* A,
+                                  const WeightType* B,
+                                  const T* weight_scales,
+                                  const T* biases,
+                                  T* C,
+                                  int64_t* total_rows_before_expert,
+                                  int64_t total_rows,
+                                  int64_t gemm_n,
+                                  int64_t gemm_k,
+                                  int num_experts,
+                                  CutlassGemmConfig gemm_config,
+                                  int sm_version,
+                                  int multi_processor_count,
+                                  cudaStream_t stream,
+                                  int* occupancy = nullptr) {
+  switch (gemm_config.tile_config) {
+    dispatch_gemm_config_macro(128, 128, 8, 64, 64, 8);
+    case CutlassTileConfig::Undefined:
+      PADDLE_FATAL(
+          "[dispatch_moe_gemm_to_cutlass][SIMT] gemm config "
+          "undefined.");
+      break;
+    case CutlassTileConfig::ChooseWithHeuristic:
+      PADDLE_FATAL(
+          "[dispatch_moe_gemm_to_cutlass][SIMT] gemm config should "
+          "have already been set by heuristic.");
+      break;
+    default:
+      PADDLE_FATAL(
+          "[dispatch_moe_gemm_to_cutlass][SIMT] Unsupported config "
+          "for float MoE gemm.");
+      break;
+  }
+}
+
+template <typename T, typename WeightType>
+MoeGemmRunner<T, WeightType>::MoeGemmRunner() {
+  int device{-1};
+  check_cuda_error(cudaGetDevice(&device));
+  sm_ = getSMVersion();
+  check_cuda_error(cudaDeviceGetAttribute(
+      &multi_processor_count_, cudaDevAttrMultiProcessorCount, device));
+}
+
+template <typename T, typename WeightType>
+template <typename EpilogueTag>
+void MoeGemmRunner<T, WeightType>::dispatch_to_arch<EpilogueTag>(
+    const T* A,
+    const WeightType* B,
+    const T* weight_scales,
+    const T* biases,
+    T* C,
+    int64_t* total_rows_before_expert,
+    int64_t total_rows,
+    int64_t gemm_n,
+    int64_t gemm_k,
+    int num_experts,
+    CutlassGemmConfig gemm_config,
+    cudaStream_t stream,
+    int* occupancy) {
+#define dispatch_moe_gemm_to_cutlass_macro(ARCH)                  \
+  dispatch_moe_gemm_to_cutlass<T, WeightType, ARCH, EpilogueTag>( \
+      A,                                                          \
+      B,                                                          \
+      weight_scales,                                              \
+      biases,                                                     \
+      C,                                                          \
+      total_rows_before_expert,                                   \
+      total_rows,                                                 \
+      gemm_n,                                                     \
+      gemm_k,                                                     \
+      num_experts,                                                \
+      gemm_config,                                                \
+      sm_,                                                        \
+      multi_processor_count_,                                     \
+      stream,                                                     \
+      occupancy);
+
+  if (sm_ >= 70 && sm_ < 75) {
+    dispatch_moe_gemm_to_cutlass_macro(cutlass::arch::Sm70);
+  } else if (sm_ >= 75 && sm_ < 80) {
+    dispatch_moe_gemm_to_cutlass_macro(cutlass::arch::Sm75);
+  } else if (sm_ >= 80 && sm_ < 91) {
+    dispatch_moe_gemm_to_cutlass_macro(cutlass::arch::Sm80);
+  } else {
+    PADDLE_FATAL("[MoE][GEMM Dispatch] Arch unsupported for MoE GEMM");
+  }
+}
+
+template <typename T, typename WeightType>
+template <typename EpilogueTag>
+void MoeGemmRunner<T, WeightType>::run_gemm<EpilogueTag>(
+    const T* A,
+    const WeightType* B,
+    const T* weight_scales,
+    const T* biases,
+    T* C,
+    int64_t* total_rows_before_expert,
+    int64_t total_rows,
+    int64_t gemm_n,
+    int64_t gemm_k,
+    int num_experts,
+    cudaStream_t stream) {
+  static constexpr bool is_weight_only = !std::is_same<T, WeightType>::value;
+  static constexpr bool only_simt_configs = std::is_same<T, float>::value;
+  std::vector<CutlassGemmConfig> candidate_configs =
+      get_candidate_configs(sm_, -1, is_weight_only, only_simt_configs, true);
+  static constexpr int warm_time = 5;
+  static constexpr int test_time = 10;
+  auto& gemmConfigManager = GemmConfigManager::Instance();
+  constexpr GemmDataType dtype = getGemmDataType<T>();
+  constexpr GemmDataType wdtype = getGemmDataType<WeightType>();
+  GemmIDType gemmId{
+      gemm_n, gemm_k, GemmType::MOEGEMM, dtype, wdtype, num_experts};
+  CutlassGemmConfig chosen_config;
+  auto chosen_config_optional =
+      gemmConfigManager.getBestConfig(gemmId, total_rows);
+  if (chosen_config_optional != std::nullopt) {
+    chosen_config = chosen_config_optional.value();
+  } else {
+    float best_time = std::numeric_limits<float>::max();
+    CutlassGemmConfig best_config;
+    int profile_total_rows =
+        std::min(gemmConfigManager.nextPowerOfTwo(total_rows),
+                 gemmConfigManager.getMaxProfileM());
+
+    for (size_t ii = 0; ii < candidate_configs.size(); ++ii) {
+      for (int i = 0; i < warm_time; i++) {
+        dispatch_to_arch<EpilogueTag>(A,
+                                      B,
+                                      weight_scales,
+                                      biases,
+                                      C,
+                                      total_rows_before_expert,
+                                      total_rows,
+                                      gemm_n,
+                                      gemm_k,
+                                      num_experts,
+                                      candidate_configs[ii],
+                                      stream);
+      }
+      cudaEvent_t start;
+      cudaEvent_t stop;
+      check_cuda_error(cudaEventCreate(&start));
+      check_cuda_error(cudaEventCreate(&stop));
+      check_cuda_error(cudaStreamSynchronize(stream));
+      check_cuda_error(cudaEventRecord(start, stream));
+      for (int i = 0; i < test_time; i++) {
+        dispatch_to_arch<EpilogueTag>(A,
+                                      B,
+                                      weight_scales,
+                                      biases,
+                                      C,
+                                      total_rows_before_expert,
+                                      total_rows,
+                                      gemm_n,
+                                      gemm_k,
+                                      num_experts,
+                                      candidate_configs[ii],
+                                      stream);
+      }
+      check_cuda_error(cudaEventRecord(stop, stream));
+      check_cuda_error(cudaEventSynchronize(stop));
+      float elapsed;
+      check_cuda_error(cudaEventElapsedTime(&elapsed, start, stop));
+      check_cuda_error(cudaEventDestroy(start));
+      check_cuda_error(cudaEventDestroy(stop));
+      if (elapsed < best_time) {
+        best_time = elapsed;
+        best_config = candidate_configs[ii];
+      }
+    }
+    gemmConfigManager.addBestConfig(gemmId, profile_total_rows, best_config);
+    chosen_config = best_config;
+  }
+
+  dispatch_to_arch<EpilogueTag>(A,
+                                B,
+                                weight_scales,
+                                biases,
+                                C,
+                                total_rows_before_expert,
+                                total_rows,
+                                gemm_n,
+                                gemm_k,
+                                num_experts,
+                                chosen_config,
+                                stream);
+}
+
+template <typename T, typename WeightType>
+void MoeGemmRunner<T, WeightType>::moe_gemm_bias_act(
+    const T* A,
+    const WeightType* B,
+    const T* weight_scales,
+    const T* biases,
+    T* C,
+    int64_t* total_rows_before_expert,
+    int64_t total_rows,
+    int64_t gemm_n,
+    int64_t gemm_k,
+    int num_experts,
+    std::string activation_type,
+    cudaStream_t stream) {
+  if (activation_type == "none") {
+    if (biases) {
+      run_gemm<EpilogueOpBias>(A,
+                               B,
+                               weight_scales,
+                               biases,
+                               C,
+                               total_rows_before_expert,
+                               total_rows,
+                               gemm_n,
+                               gemm_k,
+                               num_experts,
+                               stream);
+    } else {
+      run_gemm<EpilogueOpNoBias>(A,
+                                 B,
+                                 weight_scales,
+                                 nullptr,
+                                 C,
+                                 total_rows_before_expert,
+                                 total_rows,
+                                 gemm_n,
+                                 gemm_k,
+                                 num_experts,
+                                 stream);
+    }
+  }
+}
+
+template <typename T, typename WeightType>
+void MoeGemmRunner<T, WeightType>::moe_gemm(const T* A,
+                                            const WeightType* B,
+                                            const T* weight_scales,
+                                            T* C,
+                                            int64_t* total_rows_before_expert,
+                                            int64_t total_rows,
+                                            int64_t gemm_n,
+                                            int64_t gemm_k,
+                                            int num_experts,
+                                            cudaStream_t stream) {
+  run_gemm<EpilogueOpNoBias>(A,
+                             B,
+                             weight_scales,
+                             nullptr,
+                             C,
+                             total_rows_before_expert,
+                             total_rows,
+                             gemm_n,
+                             gemm_k,
+                             num_experts,
+                             stream);
+}

csrc/gpu/moe/fused_moe/fused_moe.cu

@@ -0,0 +1,192 @@
+// 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.
+
+
+// Ignore CUTLASS warnings about type punning
+
+#pragma once
+
+#include "cutlass/numeric_conversion.h"
+#include "helper.h"
+#include "moe/fused_moe_helper.h"
+
+__global__ void compute_total_rows_before_expert_kernel(
+    int* sorted_experts,
+    const int64_t sorted_experts_len,
+    const int64_t num_experts,
+    int64_t* total_rows_before_expert) {
+  // First, compute the global tid. We only need 1 thread per expert.
+  const int expert = blockIdx.x * blockDim.x + threadIdx.x;
+  if (expert >= num_experts) return;
+
+  // This should construct the last index where each expert occurs.
+  total_rows_before_expert[expert] =
+      find_total_elts_leq_target(sorted_experts, sorted_experts_len, expert);
+}
+
+void compute_total_rows_before_expert(int* sorted_indices,
+                                      const int64_t total_indices,
+                                      const int64_t num_experts,
+                                      int64_t* total_rows_before_expert,
+                                      cudaStream_t stream) {
+  const int threads = std::min(int64_t(1024), num_experts);
+  const int blocks = (num_experts + threads - 1) / threads;
+
+  compute_total_rows_before_expert_kernel<<<blocks, threads, 0, stream>>>(
+      sorted_indices, total_indices, num_experts, total_rows_before_expert);
+}
+
+template <paddle::DataType T>
+void FusedMoeKernel(const paddle::Tensor& input,
+                    const paddle::Tensor& gate_weight,
+                    const paddle::Tensor& ffn1_weight,
+                    const paddle::optional<paddle::Tensor>& ffn1_scale,
+                    const paddle::optional<paddle::Tensor>& ffn1_bias,
+                    const paddle::Tensor& ffn2_weight,
+                    const paddle::optional<paddle::Tensor>& ffn2_scale,
+                    const paddle::optional<paddle::Tensor>& ffn2_bias,
+                    const std::string& quant_method,
+                    const int moe_topk,
+                    const bool group_moe,
+                    const bool norm_topk_prob,
+                    paddle::Tensor* output) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  auto* output_data = output->data<data_t>();
+
+  auto fp16_moe_gemm_runner = MoeGemmRunner<DataType_, DataType_>();
+  auto int8_moe_gemm_runner = MoeGemmRunner<DataType_, uint8_t>();
+  auto int4_moe_gemm_runner = MoeGemmRunner<DataType_, cutlass::uint4b_t>();
+
+  using NvType = typename traits_::DataType;
+  auto moe_compute = MoeHelper<data_t, NvType>(quant_method,
+                                               &fp16_moe_gemm_runner,
+                                               &int8_moe_gemm_runner,
+                                               &int4_moe_gemm_runner);
+
+  moe_compute.ComputeFFN(&input,
+                         &gate_weight,
+                         &ffn1_weight,
+                         ffn1_scale ? ffn1_scale.get_ptr() : nullptr,
+                         ffn1_bias ? ffn1_bias.get_ptr() : nullptr,
+                         &ffn2_weight,
+                         ffn2_scale ? ffn2_scale.get_ptr() : nullptr,
+                         ffn2_bias ? ffn2_bias.get_ptr() : nullptr,
+                         nullptr,
+                         moe_topk,
+                         group_moe,
+                         norm_topk_prob,
+                         1.0,  // ComputeFFN
+                         "ffn",
+                         output);
+}
+
+
+std::vector<paddle::Tensor> FusedExpertMoe(
+    const paddle::Tensor& input,
+    const paddle::Tensor& gate_weight,
+    const paddle::Tensor& ffn1_weight,
+    const paddle::Tensor& ffn2_weight,
+    const paddle::optional<paddle::Tensor>& ffn1_bias,
+    const paddle::optional<paddle::Tensor>& ffn1_scale,
+    const paddle::optional<paddle::Tensor>& ffn2_bias,
+    const paddle::optional<paddle::Tensor>& ffn2_scale,
+    const std::string& quant_method,
+    const int moe_topk,
+    const bool norm_topk_prob,
+    const bool group_moe) {
+  const auto input_type = input.dtype();
+  auto output = paddle::empty_like(input);
+
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      FusedMoeKernel<paddle::DataType::BFLOAT16>(input,
+                                                 gate_weight,
+                                                 ffn1_weight,
+                                                 ffn1_scale,
+                                                 ffn1_bias,
+                                                 ffn2_weight,
+                                                 ffn2_scale,
+                                                 ffn2_bias,
+                                                 quant_method,
+                                                 moe_topk,
+                                                 group_moe,
+                                                 norm_topk_prob,
+                                                 &output);
+      break;
+    case paddle::DataType::FLOAT16:
+      FusedMoeKernel<paddle::DataType::FLOAT16>(input,
+                                                gate_weight,
+                                                ffn1_weight,
+                                                ffn1_scale,
+                                                ffn1_bias,
+                                                ffn2_weight,
+                                                ffn2_scale,
+                                                ffn2_bias,
+                                                quant_method,
+                                                moe_topk,
+                                                group_moe,
+                                                norm_topk_prob,
+                                                &output);
+      break;
+    default:
+      PD_THROW("Unsupported data type for FusedMoeKernel");
+  }
+  return {output};
+}
+
+std::vector<std::vector<int64_t>> FusedExpertMoeInferShape(
+    const std::vector<int64_t>& input_shape,
+    const std::vector<int64_t>& gate_weight_shape,
+    const std::vector<int64_t>& ffn1_weight_shape,
+    const std::vector<int64_t>& ffn2_weight_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn1_bias_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn1_scale_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn2_bias_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn2_scale_shape) {
+  return {input_shape};
+}
+
+std::vector<paddle::DataType> FusedExpertMoeInferDtype(
+    const paddle::DataType& input_dtype,
+    const paddle::DataType& gate_weight_dtype,
+    const paddle::DataType& ffn1_weight_dtype,
+    const paddle::DataType& ffn2_weight_dtype,
+    const paddle::optional<paddle::DataType>& ffn1_bias_dtype,
+    const paddle::optional<paddle::DataType>& ffn1_scale_dtype,
+    const paddle::optional<paddle::DataType>& ffn2_bias_dtype,
+    const paddle::optional<paddle::DataType>& ffn2_scale_dtype) {
+  return {input_dtype};
+}
+
+
+PD_BUILD_OP(fused_expert_moe)
+    .Inputs({"input",
+             "gate_weight",
+             "ffn1_weight",
+             "ffn2_weight",
+             paddle::Optional("ffn1_bias"),
+             paddle::Optional("ffn1_scale"),
+             paddle::Optional("ffn2_bias"),
+             paddle::Optional("ffn2_scale")})
+    .Outputs({"output"})
+    .Attrs({"quant_method:std::string",
+            "moe_topk:int",
+            "norm_topk_prob:bool",
+            "group_moe:bool"})
+    .SetKernelFn(PD_KERNEL(FusedExpertMoe))
+    .SetInferShapeFn(PD_INFER_SHAPE(FusedExpertMoeInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(FusedExpertMoeInferDtype));

csrc/gpu/moe/fused_moe/glog/logging.h

@@ -0,0 +1,14 @@
+#pragma once
+#include <iostream>
+
+/// @brief Defines a macro to disable log output.
+/// @details This macro returns a `NullStream` object, and all content passed via `<<` is ignored.
+///          It is designed to keep the code compilable without modifying log call statements.
+/// @param level The log level (ignored).
+class NullStream {
+public:
+    template <typename T>
+    NullStream& operator<<(const T&) { return *this; }
+};
+
+#define VLOG(level) NullStream()

csrc/gpu/moe/fused_moe/moe/fused_moe_helper.h

@@ -0,0 +1,427 @@
+
+/* 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. */
+
+#pragma once
+
+#include "cutlass_kernels/moe_gemm/fused_moe_gemm_kernels.h"
+#include "moe/fused_moe_op.h"
+
+using namespace phi;
+
+template <typename T, int VecSize>
+__global__ void moe_token_type_ids_kernel(T *gating_output,
+                                          const int *moe_token_type_ids_out,
+                                          const int num_rows,
+                                          const int num_experts,
+                                          const int k) {
+  const int moe_token_index = blockIdx.x * blockDim.x + threadIdx.x;
+
+  if (moe_token_index >= num_rows) {
+    return;
+  }
+
+  gating_output[moe_token_index * 2] =
+      gating_output[moe_token_index * 2] +
+      (moe_token_type_ids_out[moe_token_index]) * -1e10;
+  gating_output[moe_token_index * 2 + 1] =
+      gating_output[moe_token_index * 2 + 1] +
+      (1 - moe_token_type_ids_out[moe_token_index]) * -1e10;
+}
+
+template <typename T>
+void moe_token_type_ids_kernelLauncher(T *gating_output,
+                                       const int *moe_token_type_ids_out,
+                                       const int num_rows,
+                                       const int num_experts,
+                                       const int k,
+                                       cudaStream_t stream) {
+  const int blocks = num_rows * k / 512 + 1;
+  const int threads = 512;
+  moe_token_type_ids_kernel<T, 1><<<blocks, 512, 0, stream>>>(
+      gating_output, moe_token_type_ids_out, num_rows, num_experts, k);
+}
+
+
+template <typename T, typename NvType>
+class MoeHelper {
+public:
+  MoeHelper(const std::string gemm_method,
+            MoeGemmRunner<NvType, NvType> *fp16_moe_gemm_runner,
+            MoeGemmRunner<NvType, uint8_t> *int8_moe_gemm_runner,
+            MoeGemmRunner<NvType, cutlass::uint4b_t> *int4_moe_gemm_runner,
+            int layernum = 0)
+      : gemm_method_(gemm_method),
+        fp16_moe_gemm_runner_(fp16_moe_gemm_runner),
+        int8_moe_gemm_runner_(int8_moe_gemm_runner),
+        int4_moe_gemm_runner_(int4_moe_gemm_runner),
+        layernum_(layernum) {}
+
+  // --------      getWorkspaceSize      -------- //
+  template <typename KeyT>
+  size_t getWorkspaceSize(const int64_t num_rows,
+                          const int64_t hidden_size,
+                          const int64_t inter_size,
+                          const int64_t num_experts,
+                          const int64_t k) {
+    const size_t buf_size = AlignTo16(k * num_rows * hidden_size);
+    const size_t interbuf_size = AlignTo16(k * num_rows * inter_size);
+    const size_t padded_experts = AlignTo16(num_experts);
+    const size_t num_moe_inputs = AlignTo16(k * num_rows);
+    // softmax output, permuted_rows and permuted_experts have moved to outside
+    // of moe kernel, allocate them in Encoder or Decoder before invoking
+    // FfnLayer forward.
+    size_t total_ws_bytes =
+        5 * num_moe_inputs *
+        sizeof(int);  // source_rows_, permuted_rows_, permuted_experts_
+    total_ws_bytes += buf_size * sizeof(KeyT);  // permuted_data
+    total_ws_bytes +=
+        padded_experts * sizeof(int64_t);  // Hold total_rows_before_expert_
+
+    const size_t bytes_for_fc1_result = interbuf_size * sizeof(KeyT);
+    const size_t sorter_ws_size_bytes =
+        AlignTo16(sorter_.getWorkspaceSize(num_rows));
+    sorter_.update_num_experts(num_experts);
+
+    int64_t bytes_for_intermediate_and_sorting = bytes_for_fc1_result;
+    if (sorter_ws_size_bytes > bytes_for_fc1_result) {
+      int64_t remaining_bytes =
+          AlignTo16(sorter_ws_size_bytes - bytes_for_fc1_result);
+      bytes_for_intermediate_and_sorting += remaining_bytes;
+    }
+
+    total_ws_bytes +=
+        bytes_for_intermediate_and_sorting;  // intermediate (fc1) output + cub
+                                             // sorting workspace
+
+    int64_t num_softmax_outs = 0;
+    const bool is_pow_2 =
+        (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+    if (!is_pow_2 || num_experts > 256) {
+      num_softmax_outs = AlignTo16(num_rows * num_experts);
+    }
+
+    total_ws_bytes += num_softmax_outs * sizeof(float);
+
+    return total_ws_bytes;
+  }
+
+  void ComputeFFN(const paddle::Tensor *input,
+                  const paddle::Tensor *gate_weight,
+                  const paddle::Tensor *ffn1_weight,
+                  const paddle::Tensor *ffn1_scale,
+                  const paddle::Tensor *ffn1_bias,
+                  const paddle::Tensor *ffn2_weight,
+                  const paddle::Tensor *ffn2_scale,
+                  const paddle::Tensor *ffn2_bias,
+                  const paddle::Tensor *moe_token_type_ids,
+                  const int moe_topk,
+                  const bool group_moe,
+                  const bool norm_topk_prob,
+                  const float routed_scaling_factor,
+                  const std::string moe_type,
+                  paddle::Tensor *output) {
+    auto *input_activations = input->data<T>();
+    auto *gating_weights = gate_weight->data<float>();
+    const T *fc1_expert_biases = ffn1_bias ? ffn1_bias->data<T>() : nullptr;
+    const T *fc2_expert_biases = ffn2_bias ? ffn2_bias->data<T>() : nullptr;
+
+    auto *output_ = output->data<T>();
+    auto stream = input->stream();
+    auto place = input->place();
+    auto input_type = input->dtype();
+
+    auto input_dims = input->dims();
+    auto ffn1_dims = ffn1_weight->dims();
+    int64_t token_num = 0;
+    if (input_dims.size() == 3) {
+      token_num = input_dims[0] * input_dims[1];
+    } else {
+      token_num = input_dims[0];
+    }
+    const int64_t num_rows = token_num;
+
+    const int64_t hidden_size = ffn1_dims[1];
+    int64_t inter_dim = 0;
+    if (moe_type == "qkv") {
+      inter_dim = ffn1_dims[2] * ffn1_dims[3] * ffn1_dims[4];
+    } else {
+      inter_dim = ffn1_dims[2];
+    }
+
+    if (gemm_method_ == "weight_only_int4") {
+      inter_dim = inter_dim * 2;
+    }
+
+    const int64_t inter_size = inter_dim;
+    const int64_t num_experts = ffn1_dims[0];
+    const int64_t k = moe_topk;
+
+
+    int64_t bytes =
+        getWorkspaceSize<T>(num_rows, hidden_size, inter_size, num_experts, k);
+
+    // Pointers
+    int *expert_for_source_row;
+    int *source_rows_;
+    int *permuted_rows_;
+    int *permuted_experts_;
+    int *expanded_source_row_to_expanded_dest_row;
+
+    T *permuted_data_;
+    int64_t *total_rows_before_expert_;
+    T *fc1_result_;
+    float *softmax_out_;
+
+    paddle::Tensor ws_ptr_tensor =
+        GetEmptyTensor({bytes}, paddle::DataType::INT8, place);
+    int8_t *ws_ptr = ws_ptr_tensor.data<int8_t>();
+
+    const int64_t buf_size = AlignTo16(k * num_rows * hidden_size);
+    const int64_t interbuf_size = AlignTo16(k * num_rows * inter_size);
+    const int64_t padded_experts = AlignTo16(num_experts);
+    const int64_t num_moe_inputs = AlignTo16(k * num_rows);
+
+    expert_for_source_row = reinterpret_cast<int *>(ws_ptr);
+    source_rows_ = expert_for_source_row + num_moe_inputs;
+    permuted_rows_ = source_rows_ + num_moe_inputs;
+    permuted_experts_ = permuted_rows_ + num_moe_inputs;
+    expanded_source_row_to_expanded_dest_row =
+        permuted_experts_ + num_moe_inputs;
+    permuted_data_ = reinterpret_cast<T *>(
+        expanded_source_row_to_expanded_dest_row + num_moe_inputs);
+    total_rows_before_expert_ =
+        reinterpret_cast<int64_t *>(permuted_data_ + buf_size);
+    fc1_result_ =
+        reinterpret_cast<T *>(total_rows_before_expert_ + padded_experts);
+
+    const bool is_pow_2 =
+        (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+    if (!is_pow_2 || num_experts > 256) {
+      softmax_out_ = reinterpret_cast<float *>(fc1_result_ + interbuf_size);
+    } else {
+      softmax_out_ = nullptr;
+    }
+
+    paddle::Tensor expert_scales_float_tensor =
+        GetEmptyTensor({num_rows, moe_topk}, paddle::DataType::FLOAT32, place);
+    float *expert_scales_float = expert_scales_float_tensor.data<float>();
+
+    float *softmax_max_prob = nullptr;
+    if (group_moe) {
+      paddle::Tensor softmax_max_prob_tensor = GetEmptyTensor(
+          {num_rows, moe_topk}, paddle::DataType::FLOAT32, place);
+      // (TODO: check fill success ?)
+      paddle::experimental::fill(softmax_max_prob_tensor, 0.f);
+      softmax_max_prob = softmax_max_prob_tensor.data<float>();
+    }
+
+    paddle::Tensor fc1_out_tensor =
+        GetEmptyTensor({num_rows * k, inter_size}, input_type, place);
+    T *fc1_out = fc1_out_tensor.data<T>();
+
+    auto input_cast_tensor =
+        paddle::experimental::cast(*input, paddle::DataType::FLOAT32);
+    auto gate_tensor =
+        paddle::experimental::matmul(input_cast_tensor, *gate_weight);
+    float *gating_output = gate_tensor.data<float>();
+
+    if (moe_token_type_ids) {
+      auto *moe_token_type_ids_out = moe_token_type_ids->data<int>();
+      moe_token_type_ids_kernelLauncher<float>(gating_output,
+                                               moe_token_type_ids_out,
+                                               num_rows,
+                                               num_experts,
+                                               k,
+                                               stream);
+    }
+
+    topk_gating_softmax_kernelLauncher<float>(gating_output,
+                                              expert_scales_float,
+                                              softmax_out_,
+                                              expert_for_source_row,
+                                              source_rows_,
+                                              softmax_max_prob,
+                                              num_rows,
+                                              num_experts,
+                                              k,
+                                              group_moe,
+                                              stream);
+
+    const int64_t sorter_ws_size_bytes =
+        AlignTo16(sorter_.getWorkspaceSize(int64_t(k * num_rows)));
+
+    sorter_.run(fc1_result_,
+                sorter_ws_size_bytes,
+                expert_for_source_row,
+                permuted_experts_,
+                source_rows_,
+                permuted_rows_,
+                k * num_rows,
+                false,
+                stream);
+
+    initialize_moe_routing_kernelLauncher(
+        input_activations,
+        permuted_data_,
+        permuted_rows_,
+        expanded_source_row_to_expanded_dest_row,
+        num_rows,
+        num_rows,
+        hidden_size,
+        k,
+        stream);
+
+    const int64_t expanded_active_expert_rows = k * num_rows;
+
+    compute_total_rows_before_expert(permuted_experts_,
+                                     expanded_active_expert_rows,
+                                     num_experts,
+                                     total_rows_before_expert_,
+                                     stream);
+
+    if (gemm_method_ == "weight_only_int8") {
+      int8_moe_gemm_runner_->moe_gemm_bias_act(
+          reinterpret_cast<NvType *>(permuted_data_),
+          reinterpret_cast<const uint8_t *>(ffn1_weight->data<int8_t>()),
+          reinterpret_cast<const NvType *>(ffn1_scale->data<T>()),
+          reinterpret_cast<const NvType *>(fc1_expert_biases),
+          reinterpret_cast<NvType *>(fc1_out),
+          total_rows_before_expert_,
+          expanded_active_expert_rows,
+          inter_size,
+          hidden_size,
+          num_experts,
+          "none",
+          stream);
+    } else if (gemm_method_ == "weight_only_int4") {
+      int4_moe_gemm_runner_->moe_gemm_bias_act(
+          reinterpret_cast<NvType *>(permuted_data_),
+          reinterpret_cast<const cutlass::uint4b_t *>(
+              ffn1_weight->data<int8_t>()),
+          reinterpret_cast<const NvType *>(ffn1_scale->data<T>()),
+          reinterpret_cast<const NvType *>(fc1_expert_biases),
+          reinterpret_cast<NvType *>(fc1_out),
+          total_rows_before_expert_,
+          expanded_active_expert_rows,
+          inter_size,
+          hidden_size,
+          num_experts,
+          "none",
+          stream);
+    } else {
+      fp16_moe_gemm_runner_->moe_gemm_bias_act(
+          reinterpret_cast<NvType *>(permuted_data_),
+          reinterpret_cast<const NvType *>(ffn1_weight->data<T>()),
+          nullptr,
+          reinterpret_cast<const NvType *>(fc1_expert_biases),
+          reinterpret_cast<NvType *>(fc1_out),
+          total_rows_before_expert_,
+          expanded_active_expert_rows,
+          inter_size,
+          hidden_size,
+          num_experts,
+          "none",
+          stream);
+    }
+
+    if (moe_type == "ffn") {
+      auto act_out_tensor =
+          paddle::experimental::swiglu(fc1_out_tensor, nullptr);
+      auto act_out = act_out_tensor.data<T>();
+
+      paddle::Tensor fc2_output_tensor =
+          GetEmptyTensor({k * num_rows, hidden_size}, input_type, place);
+      T *fc2_result = fc2_output_tensor.data<T>();
+
+      if (gemm_method_ == "weight_only_int8") {
+        int8_moe_gemm_runner_->moe_gemm(
+            reinterpret_cast<NvType *>(act_out),
+            reinterpret_cast<const uint8_t *>(ffn2_weight->data<int8_t>()),
+            reinterpret_cast<const NvType *>(ffn2_scale->data<T>()),
+            reinterpret_cast<NvType *>(fc2_result),
+            total_rows_before_expert_,
+            expanded_active_expert_rows,
+            hidden_size,
+            inter_size / 2,
+            num_experts,
+            stream);
+      } else if (gemm_method_ == "weight_only_int4") {
+        int4_moe_gemm_runner_->moe_gemm(
+            reinterpret_cast<NvType *>(act_out),
+            reinterpret_cast<const cutlass::uint4b_t *>(
+                ffn2_weight->data<int8_t>()),
+            reinterpret_cast<const NvType *>(ffn2_scale->data<T>()),
+            reinterpret_cast<NvType *>(fc2_result),
+            total_rows_before_expert_,
+            expanded_active_expert_rows,
+            hidden_size,
+            inter_size / 2,
+            num_experts,
+            stream);
+      } else {
+        fp16_moe_gemm_runner_->moe_gemm(
+            reinterpret_cast<NvType *>(act_out),
+            reinterpret_cast<const NvType *>(ffn2_weight->data<T>()),
+            nullptr,
+            reinterpret_cast<NvType *>(fc2_result),
+            total_rows_before_expert_,
+            expanded_active_expert_rows,
+            hidden_size,
+            inter_size / 2,
+            num_experts,
+            stream);
+      }
+
+      finalize_moe_routing_kernelLauncher(
+          fc2_result,
+          output_,
+          fc2_expert_biases,
+          reinterpret_cast<float *>(expert_scales_float),
+          expanded_source_row_to_expanded_dest_row,
+          expert_for_source_row,
+          num_rows,
+          hidden_size,
+          k,
+          static_cast<int>(1),
+          norm_topk_prob,
+          routed_scaling_factor,
+          stream);
+    } else {
+      finalize_moe_routing_kernelLauncher(
+          // fc2_result,
+          fc1_out,
+          output_,
+          fc1_expert_biases,  // fc2_expert_biases,
+          reinterpret_cast<float *>(expert_scales_float),
+          expanded_source_row_to_expanded_dest_row,
+          expert_for_source_row,
+          num_rows,
+          inter_size,
+          k,
+          static_cast<int>(0),
+          norm_topk_prob,
+          routed_scaling_factor,
+          stream);
+    }
+  }
+
+private:
+  std::string gemm_method_;
+  MoeGemmRunner<NvType, NvType> *fp16_moe_gemm_runner_;
+  MoeGemmRunner<NvType, uint8_t> *int8_moe_gemm_runner_;
+  MoeGemmRunner<NvType, cutlass::uint4b_t> *int4_moe_gemm_runner_;
+  int layernum_;
+  CubKeyValueSorter sorter_;
+};

csrc/gpu/moe/fused_moe/moe/fused_moe_imp_op.h

@@ -0,0 +1,123 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
+ * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+ *
+ * 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 <string>
+#include <sstream>
+#include "cub/cub.cuh"
+
+static const float HALF_FLT_MAX = 65504.F;
+static const float HALF_FLT_MIN = -65504.F;
+static inline size_t AlignTo16(const size_t& input) {
+  static constexpr int ALIGNMENT = 16;
+  return ALIGNMENT * ((input + ALIGNMENT - 1) / ALIGNMENT);
+}
+
+class CubKeyValueSorter {
+ public:
+  CubKeyValueSorter() : num_experts_(0), num_bits_(sizeof(int) * 8) {}
+
+  explicit CubKeyValueSorter(const int num_experts)
+      : num_experts_(num_experts),
+        num_bits_(static_cast<int>(log2(num_experts)) + 1) {}
+
+  void update_num_experts(const int num_experts) {
+    num_experts_ = num_experts;
+    num_bits_ = static_cast<int>(log2(num_experts)) + 1;
+  }
+
+  size_t getWorkspaceSize(const size_t num_key_value_pairs,
+                          bool descending = false) {
+    num_key_value_pairs_ = num_key_value_pairs;
+    size_t required_storage = 0;
+    int* null_int = nullptr;
+    if (descending) {
+      cub::DeviceRadixSort::SortPairsDescending(NULL,
+                                                required_storage,
+                                                null_int,
+                                                null_int,
+                                                null_int,
+                                                null_int,
+                                                num_key_value_pairs,
+                                                0,
+                                                32);
+    } else {
+      cub::DeviceRadixSort::SortPairs(NULL,
+                                      required_storage,
+                                      null_int,
+                                      null_int,
+                                      null_int,
+                                      null_int,
+                                      num_key_value_pairs,
+                                      0,
+                                      num_bits_);
+    }
+    return required_storage;
+  }
+
+  template <typename KeyT>
+  void run(void* workspace,
+           const size_t workspace_size,
+           const KeyT* keys_in,
+           KeyT* keys_out,
+           const int* values_in,
+           int* values_out,
+           const size_t num_key_value_pairs,
+           bool descending,
+           cudaStream_t stream) {
+    size_t expected_ws_size = getWorkspaceSize(num_key_value_pairs);
+    size_t actual_ws_size = workspace_size;
+
+    if (expected_ws_size > workspace_size) {
+      std::stringstream err_ss;
+      err_ss << "[Error][CubKeyValueSorter::run]\n";
+      err_ss << "Error. The allocated workspace is too small to run this "
+                "problem.\n";
+      err_ss << "Expected workspace size of at least " << expected_ws_size
+             << " but got problem size " << workspace_size << "\n";
+      throw std::runtime_error(err_ss.str());
+    }
+    if (descending) {
+      cub::DeviceRadixSort::SortPairsDescending(workspace,
+                                                actual_ws_size,
+                                                keys_in,
+                                                keys_out,
+                                                values_in,
+                                                values_out,
+                                                num_key_value_pairs,
+                                                0,
+                                                32,
+                                                stream);
+    } else {
+      cub::DeviceRadixSort::SortPairs(workspace,
+                                      actual_ws_size,
+                                      keys_in,
+                                      keys_out,
+                                      values_in,
+                                      values_out,
+                                      num_key_value_pairs,
+                                      0,
+                                      num_bits_,
+                                      stream);
+    }
+  }
+
+ private:
+  size_t num_key_value_pairs_;
+  int num_experts_;
+  int num_bits_;
+};

csrc/gpu/moe/fused_moe/moe/fused_moe_op.h

@@ -0,0 +1,990 @@
+// /*
+//  * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
+//  * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
+//  *
+//  * 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_fp16.h>
+#include "moe/fused_moe_imp_op.h"
+#include "moe/fused_moe_helper.h"
+#include "cutlass/numeric_conversion.h"
+// Ignore CUTLASS warnings about type punning
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#pragma GCC diagnostic ignored "-Wunused-function"
+
+// #include "paddle/phi/backends/gpu/gpu_info.h"
+#pragma GCC diagnostic pop
+
+#include "helper.h"
+
+#define WARP_SIZE 32
+
+struct GpuLaunchConfig {
+    dim3 block_per_grid;
+    dim3 thread_per_block;
+};
+
+inline GpuLaunchConfig Get1DBlocksAnd2DGridsMoe(const int64_t cols) {
+  int blocks_x = cols;
+  int blocks_y = 1;
+  int blocks_z = 1;
+  if (blocks_x > 1024) {
+    blocks_y = 256;
+    blocks_x = (blocks_x + blocks_y - 1) / blocks_y;
+  }
+
+  GpuLaunchConfig config;
+  config.block_per_grid.x = blocks_x;
+  config.block_per_grid.y = blocks_y;
+  config.block_per_grid.z = blocks_z;
+  return config;
+}
+
+// ====================== Softmax things ===============================
+// We have our own implementation of softmax here so we can support transposing
+// the output in the softmax kernel when we extend this module to support
+// expert-choice routing.
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__
+    void group_moe_softmax(const T* input,
+                           T* output,
+                           T* softmax_max_prob,
+                           const int64_t num_cols,
+                           const int64_t softmax_num_rows) {
+  using BlockReduce = cub::BlockReduce<float, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  __shared__ float normalizing_factor;
+  __shared__ float float_max;
+  __shared__ float max_out;
+
+  int globalIdx = blockIdx.x + blockIdx.y * gridDim.x;
+  if (globalIdx >= softmax_num_rows) {
+    return;
+  }
+  const int64_t thread_row_offset = globalIdx * num_cols;
+
+  cub::Sum sum;
+  float threadData(-FLT_MAX);
+
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData = max(static_cast<float>(input[idx]), threadData);
+  }
+
+  const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
+  if (threadIdx.x == 0) {
+    float_max = maxElem;
+  }
+  __syncthreads();
+
+  threadData = 0;
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData += exp((static_cast<float>(input[idx]) - float_max));
+  }
+
+  const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum);
+
+  if (threadIdx.x == 0) {
+    normalizing_factor = 1.f / Z;
+  }
+  __syncthreads();
+
+  threadData = 0;
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    const float val =
+        exp((static_cast<float>(input[idx]) - float_max)) * normalizing_factor;
+    output[idx] = T(val);
+    threadData = max(static_cast<float>(T(val)), threadData);
+  }
+
+  const float maxOut = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
+  if (threadIdx.x == 0) {
+    // group max probs
+    max_out = 1.f / maxOut;
+    softmax_max_prob[globalIdx] = T(max_out);
+  }
+  __syncthreads();
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    // group softmax normalization
+    output[idx] = output[idx] * static_cast<T>(max_out);
+  }
+}
+
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__ void moe_top_k(const T* inputs_after_softmax,
+                                                 T* output,
+                                                 int* indices,
+                                                 int* source_rows,
+                                                 T* softmax_max_prob,
+                                                 const int64_t num_experts,
+                                                 const int64_t k,
+                                                 const int64_t num_rows) {
+  using cub_kvp = cub::KeyValuePair<int, T>;
+  using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  cub_kvp thread_kvp;
+  cub::ArgMax arg_max;
+
+  const int block_row = blockIdx.x + blockIdx.y * gridDim.x;
+  if (block_row >= num_rows) {
+    return;
+  }
+
+  const bool should_process_row = true;
+  const int thread_read_offset = block_row * num_experts;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    thread_kvp.key = 0;
+    thread_kvp.value = T(-1.f);  // This is OK because inputs are probabilities
+
+    cub_kvp inp_kvp;
+    for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
+      const int idx = thread_read_offset + expert;
+      inp_kvp.key = expert;
+      inp_kvp.value = inputs_after_softmax[idx];
+
+      for (int prior_k = 0; prior_k < k_idx; ++prior_k) {
+        const int prior_winning_expert = indices[k * block_row + prior_k];
+
+        if (prior_winning_expert == expert) {
+          inp_kvp = thread_kvp;
+        }
+      }
+
+      thread_kvp = arg_max(inp_kvp, thread_kvp);
+    }
+
+    const cub_kvp result_kvp =
+        BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+    if (threadIdx.x == 0) {
+      const int idx = k * block_row + k_idx;
+      // restore normalized probes
+      output[idx] = result_kvp.value / T(softmax_max_prob[idx]);
+      indices[idx] = should_process_row ? result_kvp.key : num_experts;
+      source_rows[idx] = k_idx * num_rows + block_row;
+    }
+    __syncthreads();
+  }
+}
+
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__ void moe_softmax(const T* input,
+                                                   T* output,
+                                                   const int64_t num_cols,
+                                                   const int64_t num_rows) {
+  using BlockReduce = cub::BlockReduce<float, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  __shared__ float normalizing_factor;
+  __shared__ float float_max;
+
+  int globalIdx = blockIdx.x + blockIdx.y * gridDim.x;
+  if (globalIdx >= num_rows) {
+    return;
+  }
+  const int64_t thread_row_offset = globalIdx * num_cols;
+
+  cub::Sum sum;
+  float threadData(-FLT_MAX);
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData = max(static_cast<float>(input[idx]), threadData);
+  }
+
+  const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
+  if (threadIdx.x == 0) {
+    float_max = maxElem;
+  }
+  __syncthreads();
+
+  threadData = 0;
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData += exp((static_cast<float>(input[idx]) - float_max));
+  }
+
+  const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum);
+
+  if (threadIdx.x == 0) {
+    normalizing_factor = 1.f / Z;
+  }
+  __syncthreads();
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    const float val =
+        exp((static_cast<float>(input[idx]) - float_max)) * normalizing_factor;
+    output[idx] = T(val);
+  }
+}
+
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__ void moe_top_k(const T* inputs_after_softmax,
+                                                 T* output,
+                                                 int* indices,
+                                                 int* source_rows,
+                                                 const int64_t num_experts,
+                                                 const int64_t k,
+                                                 const int64_t num_rows) {
+  using cub_kvp = cub::KeyValuePair<int, T>;
+  using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  cub_kvp thread_kvp;
+  cub::ArgMax arg_max;
+
+  const int block_row = blockIdx.x + blockIdx.y * gridDim.x;
+  if (block_row >= num_rows) {
+    return;
+  }
+
+  const bool should_process_row = true;
+  const int thread_read_offset = block_row * num_experts;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    thread_kvp.key = 0;
+    thread_kvp.value = T(-1.f);  // This is OK because inputs are probabilities
+
+    cub_kvp inp_kvp;
+    for (int expert = threadIdx.x; expert < num_experts; expert += TPB) {
+      const int idx = thread_read_offset + expert;
+      inp_kvp.key = expert;
+      inp_kvp.value = inputs_after_softmax[idx];
+
+      for (int prior_k = 0; prior_k < k_idx; ++prior_k) {
+        const int prior_winning_expert = indices[k * block_row + prior_k];
+
+        if (prior_winning_expert == expert) {
+          inp_kvp = thread_kvp;
+        }
+      }
+
+      thread_kvp = arg_max(inp_kvp, thread_kvp);
+    }
+
+    const cub_kvp result_kvp =
+        BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+    if (threadIdx.x == 0) {
+      const int idx = k * block_row + k_idx;
+      output[idx] = result_kvp.value;
+      indices[idx] = should_process_row ? result_kvp.key : num_experts;
+      source_rows[idx] = k_idx * num_rows + block_row;
+    }
+    __syncthreads();
+  }
+}
+
+// ====================== TopK softmax things ===============================
+
+/*
+  A Top-K gating softmax written to exploit when the number of experts in the
+  MoE layers are a small power of 2. This allows us to cleanly share the rows
+  among the threads in a single warp and eliminate communication between warps
+  (so no need to use shared mem).
+
+  It fuses the softmax, max and argmax into a single kernel.
+
+  Limitations:
+  1) This implementation is intended for when the number of experts is a small
+  power of 2. 2) This implementation assumes k is small, but will work for any
+  k.
+*/
+
+template <typename T,
+          int VPT,
+          int NUM_EXPERTS,
+          int WARPS_PER_CTA,
+          int BYTES_PER_LDG>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
+    void topk_gating_softmax(const T* input,
+                             T* output,
+                             const int64_t num_rows,
+                             int* indices,
+                             int* source_rows,
+                             const int64_t k) {
+  // We begin by enforcing compile time assertions and setting up compile time
+  // constants.
+  static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
+  static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS),
+                "NUM_EXPERTS must be power of 2");
+  static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG),
+                "BYTES_PER_LDG must be power of 2");
+  static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
+
+  // Number of bytes each thread pulls in per load
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+  static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+  static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+
+  // Restrictions based on previous section.
+  static_assert(
+      VPT % ELTS_PER_LDG == 0,
+      "The elements per thread must be a multiple of the elements per ldg");
+  static_assert(WARP_SIZE % THREADS_PER_ROW == 0,
+                "The threads per row must cleanly divide the threads per warp");
+  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW),
+                "THREADS_PER_ROW must be power of 2");
+  static_assert(THREADS_PER_ROW <= WARP_SIZE,
+                "THREADS_PER_ROW can be at most warp size");
+
+  // We have NUM_EXPERTS elements per row. We specialize for small #experts
+  static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+  static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+  static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+
+  // Restrictions for previous section.
+  static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0,
+                "The elts per row must cleanly divide the total elt per warp");
+
+  // ===================== From this point, we finally start computing run-time
+  // variables. ========================
+
+  // Compute CTA and warp rows. We pack multiple rows into a single warp, and a
+  // block contains WARPS_PER_CTA warps. This, each block processes a chunk of
+  // rows. We start by computing the start row for each block.
+  const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+
+  // Now, using the base row per thread block, we compute the base row per warp.
+  const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+
+  // The threads in a warp are split into sub-groups that will work on a row.
+  // We compute row offset for each thread sub-group
+  const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+  const int thread_row = warp_base_row + thread_row_in_warp;
+
+  // Threads with indices out of bounds should early exit here.
+  if (thread_row >= num_rows) return;
+  const bool should_process_row = true;
+
+  // We finally start setting up the read pointers for each thread. First, each
+  // thread jumps to the start of the row it will read.
+  const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+
+  // Now, we compute the group each thread belong to in order to determine the
+  // first column to start loads.
+  const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+  const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+  const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
+
+  // Determine the pointer type to use to read in the data depending on the
+  // BYTES_PER_LDG template param. In theory, this can support all powers of 2
+  // up to 16.
+  using AccessType = cutlass::AlignedArray<T, ELTS_PER_LDG>;
+
+  // Finally, we pull in the data from global mem
+  cutlass::Array<T, VPT> row_chunk_input;
+  AccessType* row_chunk_vec_ptr =
+      reinterpret_cast<AccessType*>(&row_chunk_input);
+  const AccessType* vec_thread_read_ptr =
+      reinterpret_cast<const AccessType*>(thread_read_ptr);
+#pragma unroll
+  for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
+    row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+  }
+
+  using ComputeType = float;
+  using Converter = cutlass::NumericArrayConverter<ComputeType, T, VPT>;
+  Converter compute_type_converter;
+  cutlass::Array<ComputeType, VPT> row_chunk =
+      compute_type_converter(row_chunk_input);
+
+  // First, we perform a max reduce within the thread. We can do the max in fp16
+  // safely (I think) and just convert to float afterwards for the exp + sum
+  // reduction.
+  ComputeType thread_max = row_chunk[0];
+#pragma unroll
+  for (int ii = 1; ii < VPT; ++ii) {
+    thread_max = max(thread_max, row_chunk[ii]);
+  }
+
+// Now, we find the max within the thread group and distribute among the
+// threads. We use a butterfly reduce.
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    thread_max =
+        max(thread_max,
+            __shfl_xor_sync(0xFFFFFFFF, thread_max, mask, THREADS_PER_ROW));
+  }
+
+  // From this point, thread max in all the threads have the max within the row.
+  // Now, we subtract the max from each element in the thread and take the exp.
+  // We also compute the thread local sum.
+  float row_sum = 0;
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = expf(row_chunk[ii] - thread_max);
+    row_sum += row_chunk[ii];
+  }
+
+// Now, we perform the sum reduce within each thread group. Similar to the max
+// reduce, we use a bufferfly pattern.
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    row_sum += __shfl_xor_sync(0xFFFFFFFF, row_sum, mask, THREADS_PER_ROW);
+  }
+
+  // From this point, all threads have the max and the sum for their rows in the
+  // thread_max and thread_sum variables respectively. Finally, we can scale the
+  // rows for the softmax. Technically, for top-k gating we don't need to
+  // compute the entire softmax row. We can likely look at the maxes and only
+  // compute for the top-k values in the row. However, this kernel will likely
+  // not be a bottle neck and it seems better to closer match torch and find the
+  // argmax after computing the softmax.
+  const float reciprocal_row_sum = 1.f / row_sum;
+
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
+  }
+
+  // Now, softmax_res contains the softmax of the row chunk. Now, I want to find
+  // the topk elements in each row, along with the max index.​
+  int start_col = first_elt_read_by_thread;
+  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    // First, each thread does the local argmax
+    float max_val = row_chunk[0];
+    int expert = start_col;
+#pragma unroll
+    for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD;
+         ++ldg, col += COLS_PER_GROUP_LDG) {
+#pragma unroll
+      for (int ii = 0; ii < ELTS_PER_LDG; ++ii) {
+        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+
+        // No check on the experts here since columns with the smallest index
+        // are processed first and only updated if > (not >=)
+        if (val > max_val) {
+          max_val = val;
+          expert = col + ii;
+        }
+      }
+    }
+
+// Now, we perform the argmax reduce. We use the butterfly pattern so threads
+// reach consensus about the max. This will be useful for K > 1 so that the
+// threads can agree on "who" had the max value. That thread can then blank out
+// their max with -inf and the warp can run more iterations...
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+      float other_max =
+          __shfl_xor_sync(0xFFFFFFFF, max_val, mask, THREADS_PER_ROW);
+      int other_expert =
+          __shfl_xor_sync(0xFFFFFFFF, expert, mask, THREADS_PER_ROW);
+
+      // We want lower indices to "win" in every thread so we break ties this
+      // way
+      if (other_max > max_val ||
+          (other_max == max_val && other_expert < expert)) {
+        max_val = other_max;
+        expert = other_expert;
+      }
+    }
+
+    // Write the max for this k iteration to global memory.
+    if (thread_group_idx == 0) {
+      // The lead thread from each sub-group will write out the final results to
+      // global memory. (This will be a single) thread per row of the
+      // input/output matrices.
+      const int idx = k * thread_row + k_idx;
+      output[idx] = T(max_val);
+      indices[idx] = should_process_row ? expert : NUM_EXPERTS;
+      source_rows[idx] = k_idx * num_rows + thread_row;
+    }
+
+    // Finally, we clear the value in the thread with the current max if there
+    // is another iteration to run.
+    if (k_idx + 1 < k) {
+      const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+      const int thread_to_clear_in_group =
+          (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+
+      // Only the thread in the group which produced the max will reset the
+      // "winning" value to -inf.
+      if (thread_group_idx == thread_to_clear_in_group) {
+        const int offset_for_expert = expert % ELTS_PER_LDG;
+        // Safe to set to any negative value since row_chunk values must be
+        // between 0 and 1.
+        row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] =
+            ComputeType(-10000.f);
+      }
+    }
+  }
+}
+
+namespace detail {
+// Constructs some constants needed to partition the work across threads at
+// compile time.
+template <typename T, int EXPERTS, int BYTES_PER_LDG>
+struct TopkConstants {
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 ||
+                    EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0,
+                "");
+  static constexpr int VECs_PER_THREAD =
+      std::max(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
+  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+  static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
+};
+}  // namespace detail
+
+template <typename T, int EXPERTS, int WARPS_PER_TB>
+void topk_gating_softmax_launcher_helper(const T* input,
+                                         T* output,
+                                         int* indices,
+                                         int* source_row,
+                                         const int64_t num_rows,
+                                         const int64_t num_experts,
+                                         const int64_t k,
+                                         cudaStream_t stream) {
+  static constexpr uint64_t MAX_BYTES_PER_LDG = 16;
+  static constexpr int BYTES_PER_LDG =
+      std::min(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = detail::TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int VPT = Constants::VPT;
+  static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+  const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+  const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+
+  dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
+  topk_gating_softmax<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>
+      <<<num_blocks, block_dim, 0, stream>>>(
+          input, output, num_rows, indices, source_row, k);
+}
+
+template <typename T>
+void topk_gating_softmax_kernelLauncher(const T* input,
+                                        T* output,
+                                        T* softmax,
+                                        int* indices,
+                                        int* source_row,
+                                        T* softmax_max_prob,
+                                        const int64_t num_rows,
+                                        const int64_t num_experts,
+                                        const int64_t k,
+                                        const bool group_moe,
+                                        cudaStream_t stream,
+                                        const bool topk_only_mode = false) {
+  if (topk_only_mode) {
+    static constexpr int TPB = 256;
+    const auto config_topk = Get1DBlocksAnd2DGridsMoe(num_rows);
+    moe_top_k<T, TPB><<<config_topk.block_per_grid, TPB, 0, stream>>>(
+        input, output, indices, source_row, num_experts, k, num_rows);
+    return;
+  }
+  static constexpr int WARPS_PER_TB = 4;
+
+  #define LAUNCH_TOPK_GATING_SOFTMAX_HELPER(N)                                   \
+  case N: {                                                                    \
+    topk_gating_softmax_launcher_helper<T, N, WARPS_PER_TB>(                   \
+        input, output, indices, source_row, num_rows, num_experts, k, stream); \
+    break;                                                                     \
+  }
+  switch (num_experts) {
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(2)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(4)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(8)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(16)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(32)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(64)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(128)
+    LAUNCH_TOPK_GATING_SOFTMAX_HELPER(256)
+  
+    default: {
+      static constexpr int TPB = 256;
+      if (group_moe) {
+        const int group_experts = num_experts / k;
+        const int softmax_num_rows = num_rows * k;
+        const auto config_softmax = Get1DBlocksAnd2DGridsMoe(softmax_num_rows);
+        group_moe_softmax<T, TPB>
+            <<<config_softmax.block_per_grid, TPB, 0, stream>>>(
+                input,
+                softmax,
+                softmax_max_prob,
+                group_experts,
+                softmax_num_rows);
+        const auto config_topk = Get1DBlocksAnd2DGridsMoe(num_rows);
+        moe_top_k<T, TPB>
+            <<<config_topk.block_per_grid, TPB, 0, stream>>>(softmax,
+                                                             output,
+                                                             indices,
+                                                             source_row,
+                                                             softmax_max_prob,
+                                                             num_experts,
+                                                             k,
+                                                             num_rows);
+      } else {
+        const auto config_topk = Get1DBlocksAnd2DGridsMoe(num_rows);
+        moe_softmax<T, TPB><<<config_topk.block_per_grid, TPB, 0, stream>>>(
+            input, softmax, num_experts, num_rows);
+        moe_top_k<T, TPB>
+            <<<config_topk.block_per_grid, TPB, 0, stream>>>(softmax,
+                                                             output,
+                                                             indices,
+                                                             source_row,
+                                                             num_experts,
+                                                             k,
+                                                             num_rows);
+      }
+    }
+  }
+}
+
+// ========================== Permutation things
+// =======================================
+
+// Duplicated and permutes rows for MoE. In addition, reverse the permutation
+// map to help with finalizing routing.
+
+// "expanded_x_row" simply means that the number of values is num_rows x k. It
+// is "expanded" since we will have to duplicate some rows in the input matrix
+// to match the dimensions. Duplicates will always get routed to separate
+// experts in the end.
+
+// Note that the expanded_dest_row_to_expanded_source_row map referred to here
+// has indices in the range (0, k*rows_in_input - 1). However, it is set up so
+// that index 0, rows_in_input, 2*rows_in_input ... (k-1)*rows_in_input all map
+// to row 0 in the original matrix. Thus, to know where to read in the source
+// matrix, we simply take the modulus of the expanded index.
+
+template <typename T, int VecSize>
+__global__ void initialize_moe_routing_kernel(
+    const T* unpermuted_input,
+    T* permuted_output,
+    const int* expanded_dest_row_to_expanded_source_row,
+    int* expanded_source_row_to_expanded_dest_row,
+    const int64_t num_rows,
+    const int64_t active_rows,
+    const int64_t cols,
+    const int64_t num_rows_k) {
+  using LoadT = AlignedVector<T, VecSize>;
+  LoadT src_vec;
+
+  // Reverse permutation map.
+  // I do this so that later, we can use the source -> dest map to do the k-way
+  // reduction and unpermuting. I need the reverse map for that reduction to
+  // allow each threadblock to do 1 k-way reduce without atomics later in MoE. 1
+  // thread block will be responsible for all k summations.
+  const int expanded_dest_row = blockIdx.x + blockIdx.y * gridDim.x;
+  if (expanded_dest_row >= num_rows_k) return;
+  const int expanded_source_row =
+      expanded_dest_row_to_expanded_source_row[expanded_dest_row];
+  if (threadIdx.x == 0) {
+    expanded_source_row_to_expanded_dest_row[expanded_source_row] =
+        expanded_dest_row;
+  }
+
+  if ((blockIdx.x + blockIdx.y * gridDim.x) < active_rows) {
+    // Duplicate and permute rows
+    const int source_row = expanded_source_row % num_rows;
+
+    const T* source_row_ptr = unpermuted_input + source_row * cols;
+    T* dest_row_ptr = permuted_output + expanded_dest_row * cols;
+
+    for (int tid = threadIdx.x * VecSize; tid < cols;
+         tid += blockDim.x * VecSize) {
+      // dest_row_ptr[tid] = source_row_ptr[tid];
+      Load<T, VecSize>(&source_row_ptr[tid], &src_vec);
+      Store<T, VecSize>(src_vec, &dest_row_ptr[tid]);
+    }
+  }
+}
+
+template <typename T>
+void initialize_moe_routing_kernelLauncher(
+    const T* unpermuted_input,
+    T* permuted_output,
+    const int* expanded_dest_row_to_expanded_source_row,
+    int* expanded_source_row_to_expanded_dest_row,
+    const int64_t num_rows,
+    const int64_t active_rows,
+    const int64_t cols,
+    const int64_t k,
+    cudaStream_t stream) {
+  const int threads = std::min(cols, int64_t(1024));
+  constexpr int max_pack_size = 16 / sizeof(T);
+  const auto config_initialize = Get1DBlocksAnd2DGridsMoe(num_rows * k);
+  if (cols % max_pack_size == 0) {
+    initialize_moe_routing_kernel<T, max_pack_size>
+        <<<config_initialize.block_per_grid, threads, 0, stream>>>(
+            unpermuted_input,
+            permuted_output,
+            expanded_dest_row_to_expanded_source_row,
+            expanded_source_row_to_expanded_dest_row,
+            num_rows,
+            k * active_rows,
+            cols,
+            num_rows * k);
+  } else {
+    initialize_moe_routing_kernel<T, 1>
+        <<<config_initialize.block_per_grid, threads, 0, stream>>>(
+            unpermuted_input,
+            permuted_output,
+            expanded_dest_row_to_expanded_source_row,
+            expanded_source_row_to_expanded_dest_row,
+            num_rows,
+            k * active_rows,
+            cols,
+            num_rows * k);
+  }
+}
+
+// ============================== Infer GEMM sizes
+// =================================
+__device__ inline int find_total_elts_leq_target(int* sorted_indices,
+                                                 const int64_t arr_length,
+                                                 const int64_t target) {
+  int64_t low = 0, high = arr_length - 1, target_location = -1;
+  while (low <= high) {
+    int64_t mid = (low + high) / 2;
+
+    if (sorted_indices[mid] > target) {
+      high = mid - 1;
+    } else {
+      low = mid + 1;
+      target_location = mid;
+    }
+  }
+  return target_location + 1;
+}
+
+void compute_total_rows_before_expert(int* sorted_indices,
+                                      const int64_t total_indices,
+                                      const int64_t num_experts,
+                                      int64_t* total_rows_before_expert,
+                                      cudaStream_t stream);
+
+// Final kernel to unpermute and scale
+// This kernel unpermutes the original data, does the k-way reduction and
+// performs the final skip connection.
+template <typename T, int RESIDUAL_NUM>
+__global__ void finalize_moe_routing_kernel(
+    const T* expanded_permuted_rows,
+    T* reduced_unpermuted_output,
+    const T* bias,
+    const float* scales,
+    const int* expanded_source_row_to_expanded_dest_row,
+    const int* expert_for_source_row,
+    const int64_t cols,
+    const int64_t k,
+    const int64_t compute_bias,
+    const bool norm_topk_prob,
+    const float routed_scaling_factor,
+    const int64_t num_rows) {
+  const int original_row = blockIdx.x + blockIdx.y * gridDim.x;
+  // const int original_row = blockIdx.x;
+  // const int num_rows = gridDim.x;
+  if (original_row >= num_rows) return;
+  T* reduced_row_ptr = reduced_unpermuted_output + original_row * cols;
+
+  for (int tid = threadIdx.x; tid < cols; tid += blockDim.x) {
+    T thread_output{0.f};
+    float row_rescale{0.f};
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int expanded_original_row = original_row + k_idx * num_rows;
+      const int expanded_permuted_row =
+          expanded_source_row_to_expanded_dest_row[expanded_original_row];
+
+      const int64_t k_offset = original_row * k + k_idx;
+      const float row_scale = scales[k_offset];
+      row_rescale = row_rescale + row_scale;
+
+      const T* expanded_permuted_rows_row_ptr =
+          expanded_permuted_rows + expanded_permuted_row * cols;
+
+      const int expert_idx = expert_for_source_row[k_offset];
+      const T* bias_ptr = bias ? bias + expert_idx * cols : nullptr;
+      const T bias_value = bias_ptr ? bias_ptr[tid] : T{0.f};
+
+      thread_output =
+          static_cast<float>(thread_output) +
+          row_scale * static_cast<float>(
+                          expanded_permuted_rows_row_ptr[tid] +
+                          bias_value *
+                              static_cast<T>(static_cast<float>(compute_bias)));
+    }
+
+    thread_output = static_cast<float>(thread_output) /
+                    (norm_topk_prob ? row_rescale : 1.0f) *
+                    routed_scaling_factor;
+    reduced_row_ptr[tid] = thread_output;
+  }
+}
+
+template <typename T>
+void finalize_moe_routing_kernelLauncher(
+    const T* expanded_permuted_rows,
+    T* reduced_unpermuted_output,
+    const T* bias,
+    const float* scales,
+    const int* expanded_source_row_to_expanded_dest_row,
+    const int* expert_for_source_row,
+    const int64_t num_rows,
+    const int64_t cols,
+    const int64_t k,
+    const int64_t compute_bias,
+    const bool norm_topk_prob,
+    const float routed_scaling_factor,
+    cudaStream_t stream) {
+  const int threads = std::min(cols, int64_t(1024));
+  const auto config_final = Get1DBlocksAnd2DGridsMoe(num_rows);
+
+  finalize_moe_routing_kernel<T, 1>
+      <<<config_final.block_per_grid, threads, 0, stream>>>(
+          expanded_permuted_rows,
+          reduced_unpermuted_output,
+          bias,
+          scales,
+          expanded_source_row_to_expanded_dest_row,
+          expert_for_source_row,
+          cols,
+          k,
+          compute_bias,
+          norm_topk_prob,
+          routed_scaling_factor,
+          num_rows);
+}
+
+// ========================= TopK Softmax specializations
+// ===========================
+template void topk_gating_softmax_kernelLauncher(const float*,
+                                                 float*,
+                                                 float*,
+                                                 int*,
+                                                 int*,
+                                                 float*,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const bool,
+                                                 cudaStream_t,
+                                                 const bool);
+template void topk_gating_softmax_kernelLauncher(const half*,
+                                                 half*,
+                                                 half*,
+                                                 int*,
+                                                 int*,
+                                                 half*,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const bool,
+                                                 cudaStream_t,
+                                                 const bool);
+#ifdef PADDLE_CUDA_BF16
+template void topk_gating_softmax_kernelLauncher(const __nv_bfloat16*,
+                                                 __nv_bfloat16*,
+                                                 __nv_bfloat16*,
+                                                 int*,
+                                                 int*,
+                                                 __nv_bfloat16*,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const int64_t,
+                                                 const bool,
+                                                 cudaStream_t,
+                                                 const bool);
+#endif
+// ===================== Specializations for init routing
+// =========================
+template void initialize_moe_routing_kernelLauncher(const float*,
+                                                    float*,
+                                                    const int*,
+                                                    int*,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    cudaStream_t);
+template void initialize_moe_routing_kernelLauncher(const half*,
+                                                    half*,
+                                                    const int*,
+                                                    int*,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    cudaStream_t);
+#ifdef PADDLE_CUDA_BF16
+template void initialize_moe_routing_kernelLauncher(const __nv_bfloat16*,
+                                                    __nv_bfloat16*,
+                                                    const int*,
+                                                    int*,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    const int64_t,
+                                                    cudaStream_t);
+#endif
+// ==================== Specializations for final routing
+// ===================================
+template void finalize_moe_routing_kernelLauncher(const float*,
+                                                  float*,
+                                                  const float*,
+                                                  const float*,
+                                                  const int*,
+                                                  const int*,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const bool,
+                                                  const float,
+                                                  cudaStream_t);
+template void finalize_moe_routing_kernelLauncher(const half*,
+                                                  half*,
+                                                  const half*,
+                                                  const float*,
+                                                  const int*,
+                                                  const int*,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const bool,
+                                                  const float,
+                                                  cudaStream_t);
+#ifdef PADDLE_CUDA_BF16
+template void finalize_moe_routing_kernelLauncher(const __nv_bfloat16*,
+                                                  __nv_bfloat16*,
+                                                  const __nv_bfloat16*,
+                                                  const float*,
+                                                  const int*,
+                                                  const int*,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const int64_t,
+                                                  const bool,
+                                                  const float,
+                                                  cudaStream_t);
+#endif

csrc/gpu/moe/fused_moe/moe_dispatch.cu

@@ -0,0 +1,276 @@
+// 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.
+
+
+// Ignore CUTLASS warnings about type punning
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#pragma GCC diagnostic ignored "-Wunused-function"
+#pragma once
+
+#include "moe/fused_moe_helper.h"
+#include "moe/fused_moe_op.h"
+#pragma GCC diagnostic pop
+
+#include "helper.h"
+
+
+template <paddle::DataType T>
+void MoeDispatchKernel(const paddle::Tensor& input,
+                       const paddle::Tensor& gating_output,
+                       const int moe_topk,
+                       const bool group_moe,
+                       const bool topk_only_mode,
+                       const int num_rows,
+                       const int hidden_size,
+                       const int expert_num,
+                       paddle::Tensor* permute_input,
+                       paddle::Tensor* token_nums_per_expert,
+                       paddle::Tensor* permute_indices_per_token,
+                       paddle::Tensor* expert_scales_float,
+                       paddle::Tensor* top_k_indices) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  auto stream = input.stream();
+  auto place = input.place();
+
+  if (group_moe) {
+    // Check if expert_num is divisible by moe_topk, else throw an error
+    PADDLE_ENFORCE_EQ(expert_num % moe_topk,
+                      0,
+                      common::errors::InvalidArgument(
+                          "The number of experts (expert_num) "
+                          "must be divisible by moe_topk. "
+                          "Got expert_num = %d and moe_topk = %d.",
+                          expert_num,
+                          moe_topk));
+  }
+
+  const int num_moe_inputs = AlignTo16(num_rows * moe_topk);
+  const int bytes = num_moe_inputs * sizeof(int);
+
+  CubKeyValueSorter sorter_;
+  sorter_.update_num_experts(expert_num);
+
+  const int sorter_ws_size_bytes =
+      AlignTo16(sorter_.getWorkspaceSize(moe_topk * num_rows));
+  const int sort_tmp_in_out_size = num_moe_inputs * 2 * sizeof(int);
+
+  paddle::Tensor ws_ptr_tensor =
+      GetEmptyTensor({bytes + sorter_ws_size_bytes + sort_tmp_in_out_size},
+                     paddle::DataType::INT8,
+                     place);
+
+  int8_t* ws_ptr = ws_ptr_tensor.data<int8_t>();
+  int* source_rows_ = reinterpret_cast<int*>(ws_ptr);
+  int8_t* sorter_ws_ptr = reinterpret_cast<int8_t*>(ws_ptr + bytes);
+  int* permuted_experts_ =
+      reinterpret_cast<int*>(sorter_ws_ptr + sorter_ws_size_bytes);
+  int* permuted_rows_ = permuted_experts_ + num_moe_inputs;
+
+  int* expert_for_source_row = top_k_indices->data<int>();
+
+  float* softmax_max_prob = nullptr;
+  if (group_moe) {
+    paddle::Tensor softmax_max_prob_tensor =
+        GetEmptyTensor({num_rows, moe_topk}, paddle::DataType::FLOAT32, place);
+    // (TODO: check fill success ?)
+    paddle::experimental::fill(softmax_max_prob_tensor, 0.f);
+    softmax_max_prob = softmax_max_prob_tensor.data<float>();
+  }
+
+  float* softmax_out_;
+
+  const bool is_pow_2 =
+      (expert_num != 0) && ((expert_num & (expert_num - 1)) == 0);
+
+  paddle::Tensor softmax_buffer;
+
+  if (!is_pow_2 || expert_num > 256 || group_moe) {
+    softmax_buffer = GetEmptyTensor(
+        {num_rows * expert_num}, paddle::DataType::FLOAT32, place);
+    softmax_out_ = softmax_buffer.data<float>();
+  } else {
+    softmax_out_ = nullptr;
+  }
+
+  topk_gating_softmax_kernelLauncher<float>(gating_output.data<float>(),
+                                            expert_scales_float->data<float>(),
+                                            softmax_out_,
+                                            expert_for_source_row,
+                                            source_rows_,
+                                            softmax_max_prob,
+                                            num_rows,
+                                            expert_num,
+                                            moe_topk,
+                                            group_moe,
+                                            stream,
+                                            topk_only_mode);
+
+  sorter_.run(reinterpret_cast<void*>(sorter_ws_ptr),
+              sorter_ws_size_bytes,
+              expert_for_source_row,
+              permuted_experts_,
+              source_rows_,
+              permuted_rows_,
+              moe_topk * num_rows,
+              false,
+              stream);
+
+
+  initialize_moe_routing_kernelLauncher(
+      input.data<data_t>(),
+      permute_input->data<data_t>(),
+      permuted_rows_,
+      permute_indices_per_token->data<int32_t>(),
+      num_rows,
+      num_rows,
+      hidden_size,
+      moe_topk,
+      stream);
+
+
+  compute_total_rows_before_expert(
+      permuted_experts_,
+      moe_topk * num_rows,
+      expert_num,
+      token_nums_per_expert->data<int64_t>(),
+      stream);
+}
+
+
+std::vector<paddle::Tensor> MoeExpertDispatch(
+    const paddle::Tensor& input,
+    const paddle::Tensor& gating_output,
+    const int moe_topk,
+    const bool group_moe,
+    const bool topk_only_mode) {
+  const auto input_type = input.dtype();
+  auto place = input.place();
+  int token_rows = 0;
+  auto input_dims = input.dims();
+  auto gating_dims = gating_output.dims();
+  const int expert_num = gating_dims[gating_dims.size() - 1];
+
+  if (input_dims.size() == 3) {
+    token_rows = input_dims[0] * input_dims[1];
+  } else {
+    token_rows = input_dims[0];
+  }
+  const int num_rows = token_rows;
+  const int hidden_size = input.dims()[input_dims.size() - 1];
+
+  auto permute_input =
+      GetEmptyTensor({moe_topk * num_rows, hidden_size}, input_type, place);
+  // correspond to the weighted coefficients of the results from each expert.
+  auto expert_scales_float =
+      GetEmptyTensor({num_rows, moe_topk}, paddle::DataType::FLOAT32, place);
+  auto top_k_indices =
+      GetEmptyTensor({num_rows, moe_topk}, paddle::DataType::INT32, place);
+
+  auto token_nums_per_expert =
+      GetEmptyTensor({expert_num}, paddle::DataType::INT64, place);
+  auto permute_indices_per_token =
+      GetEmptyTensor({moe_topk, num_rows}, paddle::DataType::INT32, place);
+
+
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      MoeDispatchKernel<paddle::DataType::BFLOAT16>(input,
+                                                    gating_output,
+                                                    moe_topk,
+                                                    group_moe,
+                                                    topk_only_mode,
+                                                    num_rows,
+                                                    hidden_size,
+                                                    expert_num,
+                                                    &permute_input,
+                                                    &token_nums_per_expert,
+                                                    &permute_indices_per_token,
+                                                    &expert_scales_float,
+                                                    &top_k_indices);
+      break;
+    case paddle::DataType::FLOAT16:
+      MoeDispatchKernel<paddle::DataType::FLOAT16>(input,
+                                                   gating_output,
+                                                   moe_topk,
+                                                   group_moe,
+                                                   topk_only_mode,
+                                                   num_rows,
+                                                   hidden_size,
+                                                   expert_num,
+                                                   &permute_input,
+                                                   &token_nums_per_expert,
+                                                   &permute_indices_per_token,
+                                                   &expert_scales_float,
+                                                   &top_k_indices);
+      break;
+    default:
+      PD_THROW("Unsupported data type for MoeDispatchKernel");
+  }
+  return {permute_input,
+          token_nums_per_expert,
+          permute_indices_per_token,
+          expert_scales_float,
+          top_k_indices};
+}
+
+
+std::vector<std::vector<int64_t>> MoeExpertDispatchInferShape(
+    const std::vector<int64_t>& input_shape,
+    const std::vector<int64_t>& gating_output_shape,
+    const int moe_topk) {
+  int token_rows = -1;
+
+  if (input_shape.size() == 3) {
+    token_rows = input_shape[0] * input_shape[1];
+  } else {
+    token_rows = input_shape[0];
+  }
+  const int expert_num = gating_output_shape[gating_output_shape.size() - 1];
+  const int num_rows = token_rows;
+  const int hidden_size = input_shape[input_shape.size() - 1];
+
+  return {{moe_topk * num_rows, hidden_size},
+          {expert_num},
+          {moe_topk, num_rows},
+          {num_rows, moe_topk},
+          {num_rows, moe_topk}};
+}
+
+std::vector<paddle::DataType> MoeExpertDispatchInferDtype(
+    const paddle::DataType& input_dtype,
+    const paddle::DataType& gating_output_dtype,
+    const int moe_topk) {
+  return {input_dtype,
+          paddle::DataType::INT64,
+          paddle::DataType::INT32,
+          paddle::DataType::FLOAT32,
+          paddle::DataType::INT32};
+}
+
+
+PD_BUILD_OP(moe_expert_dispatch)
+    .Inputs({"input", "gating_output"})
+    .Outputs({"permute_input",
+              "token_nums_per_expert",
+              "permute_indices_per_token",
+              "expert_scales_float",
+              "top_k_indices"})
+    .Attrs({"moe_topk:int", "group_moe:bool", "topk_only_mode:bool"})
+    .SetKernelFn(PD_KERNEL(MoeExpertDispatch))
+    .SetInferShapeFn(PD_INFER_SHAPE(MoeExpertDispatchInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MoeExpertDispatchInferDtype));

csrc/gpu/moe/fused_moe/moe_ffn.cu

@@ -0,0 +1,232 @@
+// 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 "cutlass/numeric_conversion.h"
+#include "helper.h"
+#include "moe/fused_moe_helper.h"
+
+template <paddle::DataType T>
+void MoeFFNKernel(const paddle::Tensor& permute_input,
+                  const paddle::Tensor& token_nums_per_expert,
+                  const paddle::Tensor& ffn1_weight,
+                  const paddle::Tensor& ffn2_weight,
+                  const paddle::optional<paddle::Tensor>& ffn1_bias,
+                  const paddle::optional<paddle::Tensor>& ffn1_scale,
+                  const paddle::optional<paddle::Tensor>& ffn2_scale,
+                  const std::string& quant_method,
+                  paddle::Tensor ffn_out) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+
+  auto ffn_out_data = ffn_out.data<data_t>();
+  auto permuted_data = permute_input.data<data_t>();
+  auto place = permute_input.place();
+  auto input_type = permute_input.dtype();
+  auto stream = permute_input.stream();
+
+  auto fp16_moe_gemm_runner = MoeGemmRunner<DataType_, DataType_>();
+  auto int8_moe_gemm_runner = MoeGemmRunner<DataType_, uint8_t>();
+  auto int4_moe_gemm_runner = MoeGemmRunner<DataType_, cutlass::uint4b_t>();
+
+  const int64_t expanded_active_expert_rows = permute_input.dims()[0];
+  const int num_experts = ffn1_weight.dims()[0];
+  const int hidden_size = ffn1_weight.dims()[1];
+  int inter_dim = ffn1_weight.dims()[2];
+
+  if (quant_method == "weight_only_int4") {
+    inter_dim = inter_dim * 2;
+  }
+
+  const int64_t inter_size = inter_dim;
+
+  paddle::Tensor fc1_out_tensor = GetEmptyTensor(
+      {expanded_active_expert_rows, inter_size}, input_type, place);
+  auto fc1_out = fc1_out_tensor.data<data_t>();
+
+  using NvType = typename traits_::DataType;
+
+  auto fc1_expert_biases =
+      ffn1_bias
+          ? const_cast<paddle::Tensor*>(ffn1_bias.get_ptr())->data<data_t>()
+          : nullptr;
+
+  if (quant_method == "weight_only_int8") {
+    int8_moe_gemm_runner.moe_gemm_bias_act(
+        reinterpret_cast<const NvType*>(permuted_data),
+        reinterpret_cast<const uint8_t*>(ffn1_weight.data<int8_t>()),
+        reinterpret_cast<const NvType*>(
+            const_cast<paddle::Tensor*>(ffn1_scale.get_ptr())->data<data_t>()),
+        reinterpret_cast<const NvType*>(fc1_expert_biases),
+        reinterpret_cast<NvType*>(fc1_out),
+        const_cast<int64_t*>(token_nums_per_expert.data<int64_t>()),
+        expanded_active_expert_rows,
+        inter_size,
+        hidden_size,
+        num_experts,
+        "none",
+        stream);
+  } else if (quant_method == "weight_only_int4") {
+    int4_moe_gemm_runner.moe_gemm_bias_act(
+        reinterpret_cast<const NvType*>(permuted_data),
+        reinterpret_cast<const cutlass::uint4b_t*>(ffn1_weight.data<int8_t>()),
+        reinterpret_cast<const NvType*>(
+            const_cast<paddle::Tensor*>(ffn1_scale.get_ptr())->data<data_t>()),
+        reinterpret_cast<const NvType*>(fc1_expert_biases),
+        reinterpret_cast<NvType*>(fc1_out),
+        const_cast<int64_t*>(token_nums_per_expert.data<int64_t>()),
+        expanded_active_expert_rows,
+        inter_size,
+        hidden_size,
+        num_experts,
+        "none",
+        stream);
+  } else {
+    fp16_moe_gemm_runner.moe_gemm_bias_act(
+        reinterpret_cast<const NvType*>(permuted_data),
+        reinterpret_cast<const NvType*>(ffn1_weight.data<data_t>()),
+        nullptr,
+        reinterpret_cast<const NvType*>(fc1_expert_biases),
+        reinterpret_cast<NvType*>(fc1_out),
+        const_cast<int64_t*>(token_nums_per_expert.data<int64_t>()),
+        expanded_active_expert_rows,
+        inter_size,
+        hidden_size,
+        num_experts,
+        "none",
+        stream);
+  }
+
+  auto act_out_tensor = paddle::experimental::swiglu(fc1_out_tensor, nullptr);
+  auto act_out = act_out_tensor.data<data_t>();
+
+  if (quant_method == "weight_only_int8") {
+    int8_moe_gemm_runner.moe_gemm(
+        reinterpret_cast<const NvType*>(act_out),
+        reinterpret_cast<const uint8_t*>(ffn2_weight.data<int8_t>()),
+        reinterpret_cast<const NvType*>(
+            const_cast<paddle::Tensor*>(ffn2_scale.get_ptr())->data<data_t>()),
+        reinterpret_cast<NvType*>(ffn_out_data),
+        const_cast<int64_t*>(token_nums_per_expert.data<int64_t>()),
+        expanded_active_expert_rows,
+        hidden_size,
+        inter_size / 2,
+        num_experts,
+        stream);
+
+  } else if (quant_method == "weight_only_int4") {
+    int4_moe_gemm_runner.moe_gemm(
+        reinterpret_cast<const NvType*>(act_out),
+        reinterpret_cast<const cutlass::uint4b_t*>(ffn2_weight.data<int8_t>()),
+        reinterpret_cast<const NvType*>(
+            const_cast<paddle::Tensor*>(ffn2_scale.get_ptr())->data<data_t>()),
+        reinterpret_cast<NvType*>(ffn_out_data),
+        const_cast<int64_t*>(token_nums_per_expert.data<int64_t>()),
+        expanded_active_expert_rows,
+        hidden_size,
+        inter_size / 2,
+        num_experts,
+        stream);
+  } else {
+    fp16_moe_gemm_runner.moe_gemm(
+        reinterpret_cast<const NvType*>(act_out),
+        reinterpret_cast<const NvType*>(ffn2_weight.data<data_t>()),
+        nullptr,
+        reinterpret_cast<NvType*>(ffn_out_data),
+        const_cast<int64_t*>(token_nums_per_expert.data<int64_t>()),
+        expanded_active_expert_rows,
+        hidden_size,
+        inter_size / 2,
+        num_experts,
+        stream);
+  }
+}
+
+std::vector<paddle::Tensor> MoeExpertFFN(
+    const paddle::Tensor& permute_input,
+    const paddle::Tensor& token_nums_per_expert,
+    const paddle::Tensor& ffn1_weight,
+    const paddle::Tensor& ffn2_weight,
+    const paddle::optional<paddle::Tensor>& ffn1_bias,
+    const paddle::optional<paddle::Tensor>& ffn1_scale,
+    const paddle::optional<paddle::Tensor>& ffn2_scale,
+    const std::string& quant_method) {
+  const auto input_type = permute_input.dtype();
+  auto ffn_out = paddle::empty_like(permute_input);
+
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      MoeFFNKernel<paddle::DataType::BFLOAT16>(permute_input,
+                                               token_nums_per_expert,
+                                               ffn1_weight,
+                                               ffn2_weight,
+                                               ffn1_bias,
+                                               ffn1_scale,
+                                               ffn2_scale,
+                                               quant_method,
+                                               ffn_out);
+      break;
+    case paddle::DataType::FLOAT16:
+      MoeFFNKernel<paddle::DataType::FLOAT16>(permute_input,
+                                              token_nums_per_expert,
+                                              ffn1_weight,
+                                              ffn2_weight,
+                                              ffn1_bias,
+                                              ffn1_scale,
+                                              ffn2_scale,
+                                              quant_method,
+                                              ffn_out);
+      break;
+    default:
+      PD_THROW("Unsupported data type for MoeExpertFFN");
+  }
+  return {ffn_out};
+}
+
+std::vector<std::vector<int64_t>> MoeExpertFFNInferShape(
+    const std::vector<int64_t>& permute_input_shape,
+    const std::vector<int64_t>& token_nums_per_expert_shape,
+    const std::vector<int64_t>& ffn1_weight_shape,
+    const std::vector<int64_t>& ffn2_weight_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn1_bias_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn1_scale_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn2_scale_shape) {
+  return {permute_input_shape};
+}
+
+std::vector<paddle::DataType> MoeExpertFFNInferDtype(
+    const paddle::DataType& permute_input_dtype,
+    const paddle::DataType& token_nums_per_expert_dtype,
+    const paddle::DataType& ffn1_weight_dtype,
+    const paddle::DataType& ffn2_weight_dtype,
+    const paddle::optional<paddle::DataType>& ffn1_bias_dtype,
+    const paddle::optional<paddle::DataType>& ffn1_scale_dtype,
+    const paddle::optional<paddle::DataType>& ffn2_scale_dtype) {
+  return {permute_input_dtype};
+}
+
+PD_BUILD_OP(moe_expert_ffn)
+    .Inputs({"permute_input",
+             "token_nums_per_expert",
+             "ffn1_weight",
+             "ffn2_weight",
+             paddle::Optional("ffn1_bias"),
+             paddle::Optional("ffn1_scale"),
+             paddle::Optional("ffn2_scale")})
+    .Outputs({"output_tensor"})
+    .Attrs({"quant_method:std::string"})
+    .SetKernelFn(PD_KERNEL(MoeExpertFFN))
+    .SetInferShapeFn(PD_INFER_SHAPE(MoeExpertFFNInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MoeExpertFFNInferDtype));

csrc/gpu/moe/fused_moe/moe_reduce.cu

@@ -0,0 +1,139 @@
+// 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.
+
+
+// Ignore CUTLASS warnings about type punning
+
+
+#pragma once
+
+#include "helper.h"
+#include "moe/fused_moe_helper.h"
+#include "moe/fused_moe_op.h"
+
+template <paddle::DataType T>
+void MoeReduceKernel(const paddle::Tensor& ffn_out,
+                     const paddle::Tensor& expert_scales_float,
+                     const paddle::Tensor& permute_indices_per_token,
+                     const paddle::Tensor& top_k_indices,
+                     const paddle::optional<paddle::Tensor>& ffn2_bias,
+                     const bool norm_topk_prob,
+                     const float routed_scaling_factor,
+                     const int num_rows,
+                     const int hidden_size,
+                     const int topk,
+                     paddle::Tensor* output) {
+  typedef PDTraits<T> traits_;
+  typedef typename traits_::DataType DataType_;
+  typedef typename traits_::data_t data_t;
+  auto stream = ffn_out.stream();
+
+  finalize_moe_routing_kernelLauncher(
+      ffn_out.data<data_t>(),
+      output->data<data_t>(),
+      ffn2_bias ? ffn2_bias->data<data_t>() : nullptr,
+      expert_scales_float.data<float>(),
+      permute_indices_per_token.data<int32_t>(),
+      top_k_indices.data<int>(),
+      num_rows,
+      hidden_size,
+      topk,
+      static_cast<int>(1),
+      norm_topk_prob,
+      routed_scaling_factor,
+      stream);
+}
+
+
+std::vector<paddle::Tensor> MoeExpertReduce(
+    const paddle::Tensor& ffn_out,
+    const paddle::Tensor& expert_scales_float,
+    const paddle::Tensor& permute_indices_per_token,
+    const paddle::Tensor& top_k_indices,
+    const paddle::optional<paddle::Tensor>& ffn2_bias,
+    const bool norm_topk_prob,
+    const float routed_scaling_factor) {
+  const auto input_type = ffn_out.dtype();
+  auto place = ffn_out.place();
+
+  const int topk = top_k_indices.dims()[1];
+  const int num_rows = ffn_out.dims()[0] / topk;
+  const int hidden_size = ffn_out.dims()[1];
+
+  auto output = GetEmptyTensor({num_rows, hidden_size}, input_type, place);
+
+  switch (input_type) {
+    case paddle::DataType::BFLOAT16:
+      MoeReduceKernel<paddle::DataType::BFLOAT16>(ffn_out,
+                                                  expert_scales_float,
+                                                  permute_indices_per_token,
+                                                  top_k_indices,
+                                                  ffn2_bias,
+                                                  norm_topk_prob,
+                                                  routed_scaling_factor,
+                                                  num_rows,
+                                                  hidden_size,
+                                                  topk,
+                                                  &output);
+      break;
+    case paddle::DataType::FLOAT16:
+      MoeReduceKernel<paddle::DataType::BFLOAT16>(ffn_out,
+                                                  expert_scales_float,
+                                                  permute_indices_per_token,
+                                                  top_k_indices,
+                                                  ffn2_bias,
+                                                  norm_topk_prob,
+                                                  routed_scaling_factor,
+                                                  num_rows,
+                                                  hidden_size,
+                                                  topk,
+                                                  &output);
+      break;
+    default:
+      PD_THROW("Unsupported data type for MoeDispatchKernel");
+  }
+  return {output};
+}
+
+
+std::vector<std::vector<int64_t>> MoeExpertReduceInferShape(
+    const std::vector<int64_t>& ffn_out_shape,
+    const std::vector<int64_t>& expert_scales_float_shape,
+    const std::vector<int64_t>& permute_indices_per_token_shape,
+    const std::vector<int64_t>& top_k_indices_shape,
+    const paddle::optional<std::vector<int64_t>>& ffn2_bias_shape) {
+  return {ffn_out_shape};
+}
+
+std::vector<paddle::DataType> MoeExpertReduceInferDtype(
+    const paddle::DataType& ffn_out_dtype,
+    const paddle::DataType& expert_scales_float_dtype,
+    const paddle::DataType& permute_indices_per_token_dtype,
+    const paddle::DataType& top_k_indices_dtype,
+    const paddle::optional<paddle::DataType>& ffn2_bias_dtype) {
+  return {ffn_out_dtype};
+}
+
+
+PD_BUILD_OP(moe_expert_reduce)
+    .Inputs({"ffn_out",
+             "expert_scales_float",
+             "permute_indices_per_token",
+             "top_k_indices",
+             paddle::Optional("ffn2_bias")})
+    .Outputs({"output"})
+    .Attrs({"norm_topk_prob:bool", "routed_scaling_factor:float"})
+    .SetKernelFn(PD_KERNEL(MoeExpertReduce))
+    .SetInferShapeFn(PD_INFER_SHAPE(MoeExpertReduceInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(MoeExpertReduceInferDtype));

csrc/gpu/moe/preprocess_for_moe.cu

@@ -0,0 +1,180 @@
+// 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 "paddle/extension.h"
+
+
+#define CEILDIV(a,b) (((a+b-1)/b))
+
+template <typename scalar_t>
+__global__ void count_and_sort_expert_tokens_kernel(const scalar_t* __restrict__ topk_ids,
+                                                    int32_t* __restrict__ sorted_token_ids,
+                                                    int32_t* __restrict__ cumsum_buffer, 
+                                                    size_t numel) {
+  const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i];
+    int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
+    sorted_token_ids[rank_post_pad] = i;
+  }
+}
+
+template <typename scalar_t, int num_experts>
+__global__ void moe_align_block_size_kernel(const scalar_t* __restrict__ topk_ids,
+                                            int32_t* __restrict__ expert_ids,
+                                            int32_t* __restrict__ total_tokens_post_pad, 
+                                            int32_t block_size, 
+                                            size_t numel, 
+                                            int32_t* __restrict__ cumsum_buffer) {
+  __shared__ int32_t tokens_per_ep[num_experts];
+
+  for (int i = threadIdx.x; i < num_experts; i += blockDim.x) {
+      tokens_per_ep[i] = 0;
+  }
+
+  __syncthreads();
+
+  for (int i = threadIdx.x; i < numel; i += blockDim.x) {
+    int expert_id = topk_ids[i];
+    atomicAdd(&tokens_per_ep[expert_id], 1);
+  }
+
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    cumsum_buffer[0] = 0;
+    for (int i = 1; i <= num_experts; ++i) {
+      int expert_count = tokens_per_ep[i-1];
+      cumsum_buffer[i] = cumsum_buffer[i - 1] + CEILDIV(expert_count, block_size) * block_size;
+    }
+    *total_tokens_post_pad = cumsum_buffer[num_experts];
+  }
+
+  __syncthreads();
+
+  if (threadIdx.x < num_experts) {
+    for (int i = cumsum_buffer[threadIdx.x]; i < cumsum_buffer[threadIdx.x + 1]; i += block_size) {
+      expert_ids[i / block_size] = threadIdx.x;
+    }
+  }
+}
+
+
+std::vector<std::vector<int64_t>> preprocess_for_moeInferShape(const std::vector<int64_t>& topk_ids, int64_t num_experts, int64_t block_size) {
+
+
+    int topk_ids_numel = topk_ids[0] * topk_ids[1];
+    int max_num_tokens_padded = topk_ids_numel + num_experts * (block_size - 1);
+
+    std::vector<int64_t> sorted_ids = {max_num_tokens_padded};
+
+    int max_num_m_blocks = max_num_tokens_padded / block_size;
+    std::vector<int64_t> expert_ids = {max_num_m_blocks};
+    std::vector<int64_t> num_tokens_post_pad = {1};
+
+    return {sorted_ids, expert_ids, num_tokens_post_pad};
+}
+
+std::vector<paddle::DataType> preprocess_for_moeIferDtype(const paddle::DataType& topk_ids, int64_t num_experts, int64_t block_size) {
+    return {paddle::DataType::INT32, paddle::DataType::INT32, paddle::DataType::INT32};
+}
+
+
+
+std::vector<paddle::Tensor> preprocess_for_moe_kernel(const paddle::Tensor& topk_ids, int64_t num_experts, int64_t block_size) {
+    
+    int topk_ids_numel = topk_ids.shape()[0] * topk_ids.shape()[1];
+    int max_num_tokens_padded = topk_ids_numel + num_experts * (block_size - 1);
+
+    auto sorted_ids = paddle::full(
+        {max_num_tokens_padded}, 
+        topk_ids_numel,
+        paddle::DataType::INT32, 
+        topk_ids.place()
+    );
+
+    int max_num_m_blocks = max_num_tokens_padded / block_size;
+    
+    auto expert_ids = paddle::empty(
+        {max_num_m_blocks}, paddle::DataType::INT32, 
+        topk_ids.place()
+    );
+
+    auto num_tokens_post_pad = paddle::empty(
+        {1}, 
+        paddle::DataType::INT32, 
+         topk_ids.place()
+    );
+
+    auto cumsum_buffer = paddle::empty(
+        {num_experts + 1}, 
+        paddle::DataType::INT32, 
+        topk_ids.place()
+    );
+
+    auto stream = topk_ids.stream();
+    using scalar_t = int64_t;
+
+    # define run_align_kernel(num_experts) \
+    auto align_kernel = moe_align_block_size_kernel<scalar_t, num_experts>; \
+    align_kernel<<<1, 1024, 0, stream>>>( \
+    topk_ids.data<scalar_t>(),  \
+    expert_ids.data<int32_t>(), \
+    num_tokens_post_pad.data<int32_t>(), \
+    block_size,  \
+    topk_ids_numel, \
+    cumsum_buffer.data<int32_t>());
+    
+    if (num_experts == 8) {
+      run_align_kernel(8);
+    } else if (num_experts == 256) {
+      run_align_kernel(256);
+    } else if (num_experts == 2) {
+      run_align_kernel(2);
+    } else {
+      printf("errors");
+    }
+
+    const int block_threads = 256;
+    const int num_blocks = CEILDIV(topk_ids_numel, block_threads);
+    const int max_blocks = 65535;
+    const int actual_blocks = std::min(num_blocks, max_blocks);
+
+    auto sort_kernel = count_and_sort_expert_tokens_kernel<scalar_t>;
+
+    sort_kernel<<<actual_blocks, block_threads, 0, stream>>>(topk_ids.data<scalar_t>(),
+                                                              sorted_ids.data<int32_t>(),
+                                                              cumsum_buffer.data<int32_t>(), 
+                                                              topk_ids_numel);
+
+
+
+    return {sorted_ids, expert_ids, num_tokens_post_pad};
+}
+
+PD_BUILD_OP(preprocess_for_moe)
+    .Inputs({"topk_ids"})
+    .Attrs({"num_experts: int64_t", "block_size: int64_t"})
+    .Outputs({"sorted_ids", "expert_ids", "num_tokens_post_pad"})
+    .SetKernelFn(PD_KERNEL(preprocess_for_moe_kernel))
+    .SetInferShapeFn(PD_INFER_SHAPE(preprocess_for_moeInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(preprocess_for_moeIferDtype));
+
+
+
+
+
+

csrc/gpu/multi_head_latent_attention.cu

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "append_attn/multi_head_latent_attention_kernel.h"
+#include "mla_attn/batch_mla_with_paged_kv_cache.h"
 
 template <paddle::DataType D>
 std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
@@ -23,6 +24,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
     const paddle::Tensor& seq_lens_encoder,
     const paddle::Tensor& seq_lens_decoder,
     const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& cu_seqlens_q,
     const paddle::Tensor& padding_offsets,
     const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
@@ -35,6 +37,8 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
     const paddle::Tensor& decoder_batch_ids,
     const paddle::Tensor& decoder_tile_ids_per_batch,
     const paddle::Tensor& decoder_num_blocks,
+    const paddle::Tensor& decoder_num_blocks_cpu,
+    const paddle::Tensor& decoder_chunk_size_cpu,
     const paddle::Tensor& max_enc_len_this_time,
     const paddle::Tensor& max_dec_len_this_time,
     const paddle::Tensor& max_len_kv,
@@ -55,23 +59,22 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
     const float quant_max_bound,
     const float quant_min_bound,
     const float out_linear_in_scale,
-    const int speculate_max_draft_token_num,
+    const int speculate_draft_total_token_num,
     const bool causal,
     const bool speculate_decoder) {
   typedef PDTraits<D> traits_;
-  typedef typename traits_::DataType DataType_;
   typedef typename traits_::data_t data_t;
-
-  int encoder_num_blocks_data = encoder_num_blocks.data<int>()[0];
-  int kv_num_blocks_data = kv_num_blocks.data<int>()[0];
-  int decoder_num_blocks_data = decoder_num_blocks.data<int>()[0];
-  int max_enc_len_this_time_data = max_enc_len_this_time.data<int>()[0];
+  int decoder_num_blocks_data = decoder_num_blocks_cpu.data<int>()[0];
   int max_dec_len_this_time_data = max_dec_len_this_time.data<int>()[0];
+  int chunk_size = decoder_chunk_size_cpu.data<int>()[0];
   int max_len_kv_data = max_len_kv.data<int>()[0];
-  const int encoder_block_shape_q = get_encoder_block_shape_q();
-  const int decoder_block_shape_q = get_decoder_block_shape_q();
-  auto main_stream = query.stream();
+  auto mla_use_tensorcore = GetMlaUseTensorcore();
 
+  if (speculate_decoder && !mla_use_tensorcore) {
+    PD_THROW("Please use speculate_decoder=0 and FLAGS_mla_use_tensorcore=0 when sm < 90.");
+  }
+
+  auto main_stream = query.stream();
   paddle::Tensor fmha_out = paddle::full(
       {meta_data.token_nums, meta_data.q_num_heads * meta_data.head_dims_v},
       0,
@@ -79,25 +82,62 @@ std::vector<paddle::Tensor> MultiHeadLatentAttentionKernel(
       query.place());
 
   if (max_dec_len_this_time_data > 0) {
-    DecodeMLAAttentionKernel<data_t>(meta_data,
-                                     query,  // [token_num, num_heads, head_dim]
-                                     key_cache,
-                                     value_cache,
-                                     attn_mask,
-                                     out_linear_shifts,
-                                     out_linear_smooths,
-                                     seq_lens_this_time,  // q_seq_len is 1
-                                     seq_lens_decoder,
-                                     padding_offsets,
-                                     cum_offsets,
-                                     block_tables,
-                                     max_input_length,
-                                     max_len_kv_data,
-                                     softmax_scale,
-                                     out_linear_in_scale,
-                                     causal,
-                                     main_stream,
-                                     &fmha_out);
+    if (mla_use_tensorcore) {
+      BatchMLAWithPagedKVCacheKernel<data_t>(meta_data,
+                                             query,
+                                             key_cache,
+                                             attn_mask,
+                                             cache_k_dequant_scales,
+                                             cache_v_dequant_scales,
+                                             cache_k_zp,
+                                             cache_v_zp,
+                                             out_linear_shifts,
+                                             out_linear_smooths,
+                                             seq_lens_this_time,
+                                             seq_lens_decoder,
+                                             seq_lens_encoder,
+                                             cu_seqlens_q,
+                                             padding_offsets,
+                                             cum_offsets,
+                                             block_tables,
+                                             decoder_batch_ids,
+                                             decoder_tile_ids_per_batch,
+                                             decoder_num_blocks,
+                                             cache_quant_type_str,
+                                             decoder_num_blocks_data,
+                                             chunk_size,
+                                             max_input_length,
+                                             max_len_kv_data,
+                                             softmax_scale,
+                                             quant_max_bound,
+                                             quant_min_bound,
+                                             out_linear_in_scale,
+                                             speculate_draft_total_token_num,
+                                             causal,
+                                             main_stream,
+                                             &fmha_out);
+    } else {
+      DecodeMLAAttentionKernel<data_t>(
+          meta_data,
+          query,  // [token_num, num_heads, head_dim]
+          key_cache,
+          value_cache,
+          attn_mask,
+          out_linear_shifts,
+          out_linear_smooths,
+          seq_lens_this_time,  // q_seq_len is 1
+          seq_lens_decoder,
+          padding_offsets,
+          cum_offsets,
+          block_tables,
+          max_input_length,
+          max_len_kv_data,
+          softmax_scale,
+          out_linear_in_scale,
+          causal,
+          main_stream,
+          &fmha_out);
+    }
   }
   return {fmha_out};
 }
@@ -109,6 +149,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
     const paddle::Tensor& seq_lens_encoder,
     const paddle::Tensor& seq_lens_decoder,
     const paddle::Tensor& seq_lens_this_time,
+    const paddle::Tensor& cu_seqlens_q,
     const paddle::Tensor& padding_offsets,
     const paddle::Tensor& cum_offsets,
     const paddle::Tensor& block_tables,
@@ -121,6 +162,8 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
     const paddle::Tensor& decoder_batch_ids,
     const paddle::Tensor& decoder_tile_ids_per_batch,
     const paddle::Tensor& decoder_num_blocks,
+    const paddle::Tensor& decoder_num_blocks_cpu,
+    const paddle::Tensor& decoder_chunk_size_cpu,
     const paddle::Tensor& max_enc_len_this_time,
     const paddle::Tensor& max_dec_len_this_time,
     const paddle::Tensor& max_len_kv,
@@ -143,7 +186,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
     const float quant_max_bound,
     const float quant_min_bound,
     const float out_linear_in_scale,
-    const int speculate_max_draft_token_num,
+    const int speculate_draft_total_token_num,
     const bool causal,
     const bool speculate_decoder) {
   AppendAttnMetaData meta_data;
@@ -171,6 +214,7 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           seq_lens_encoder,
           seq_lens_decoder,
           seq_lens_this_time,
+          cu_seqlens_q,
           padding_offsets,
           cum_offsets,
           block_tables,
@@ -183,6 +227,8 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           decoder_batch_ids,
           decoder_tile_ids_per_batch,
           decoder_num_blocks,
+          decoder_num_blocks_cpu,
+          decoder_chunk_size_cpu,
           max_enc_len_this_time,
           max_dec_len_this_time,
           max_len_kv,
@@ -203,7 +249,55 @@ std::vector<paddle::Tensor> MultiHeadLatentAttention(
           quant_max_bound,
           quant_min_bound,
           out_linear_in_scale,
-          speculate_max_draft_token_num,
+          speculate_draft_total_token_num,
+          causal,
+          speculate_decoder);
+    }
+    case paddle::DataType::FLOAT16: {
+      return MultiHeadLatentAttentionKernel<paddle::DataType::FLOAT16>(
+          meta_data,
+          query,
+          key_cache,
+          value_cache,
+          seq_lens_encoder,
+          seq_lens_decoder,
+          seq_lens_this_time,
+          cu_seqlens_q,
+          padding_offsets,
+          cum_offsets,
+          block_tables,
+          encoder_batch_ids,
+          encoder_tile_ids_per_batch,
+          encoder_num_blocks,
+          kv_batch_ids,
+          kv_tile_ids_per_batch,
+          kv_num_blocks,
+          decoder_batch_ids,
+          decoder_tile_ids_per_batch,
+          decoder_num_blocks,
+          decoder_num_blocks_cpu,
+          decoder_chunk_size_cpu,
+          max_enc_len_this_time,
+          max_dec_len_this_time,
+          max_len_kv,
+          attn_mask,
+          query_bias,
+          query_out_scales,
+          cache_k_quant_scales,
+          cache_v_quant_scales,
+          cache_k_dequant_scales,
+          cache_v_dequant_scales,
+          cache_k_zp,
+          cache_v_zp,
+          out_linear_shifts,
+          out_linear_smooths,
+          cache_quant_type_str,
+          max_input_length,
+          softmax_scale,
+          quant_max_bound,
+          quant_min_bound,
+          out_linear_in_scale,
+          speculate_draft_total_token_num,
           causal,
           speculate_decoder);
     }
@@ -223,6 +317,7 @@ std::vector<std::vector<int64_t>> MultiHeadLatentAttentionInferShape(
     const std::vector<int64_t>& seq_lens_encoder_shape,
     const std::vector<int64_t>& seq_lens_decoder_shape,
     const std::vector<int64_t>& seq_lens_this_time_shape,
+    const std::vector<int64_t>& cu_seqlens_q_shape,
     const std::vector<int64_t>& padding_offsets_shape,
     const std::vector<int64_t>& cum_offsets_shape,
     const std::vector<int64_t>& block_tables_shape,
@@ -235,6 +330,8 @@ std::vector<std::vector<int64_t>> MultiHeadLatentAttentionInferShape(
     const std::vector<int64_t>& decoder_batch_ids_shape,
     const std::vector<int64_t>& decoder_tile_ids_per_batch_shape,
     const std::vector<int64_t>& decoder_num_blocks_shape,
+    const std::vector<int64_t>& decoder_num_blocks_cpu_shape,
+    const std::vector<int64_t>& decoder_chunk_size_cpu_shape,
     const std::vector<int64_t>& max_enc_len_this_time_shape,
     const std::vector<int64_t>& max_dec_len_this_time_shape,
     const std::vector<int64_t>& max_len_kv_shape,
@@ -257,7 +354,7 @@ std::vector<std::vector<int64_t>> MultiHeadLatentAttentionInferShape(
     const float quant_max_bound,
     const float quant_min_bound,
     const float out_linear_in_scale,
-    const int speculate_max_draft_token_num,
+    const int speculate_draft_total_token_num,
     const bool causal,
     const bool speculate_decoder) {
   const int token_num = query_shape[0];
@@ -269,6 +366,8 @@ std::vector<std::vector<int64_t>> MultiHeadLatentAttentionInferShape(
   return {{token_num, num_heads * head_dim_v}};
 }
 
+
+
 std::vector<paddle::DataType> MultiHeadLatentAttentionInferDtype(
     const paddle::DataType& query_dtype,
     const paddle::DataType& key_cache_dtype,
@@ -276,6 +375,7 @@ std::vector<paddle::DataType> MultiHeadLatentAttentionInferDtype(
     const paddle::DataType& seq_lens_encoder_dtype,
     const paddle::DataType& seq_lens_decoder_dtype,
     const paddle::DataType& seq_lens_this_time_dtype,
+    const paddle::DataType& cu_seqlens_q_dtype,
     const paddle::DataType& padding_offsets_dtype,
     const paddle::DataType& cum_offsets_dtype,
     const paddle::DataType& block_tables_dtype,
@@ -288,6 +388,8 @@ std::vector<paddle::DataType> MultiHeadLatentAttentionInferDtype(
     const paddle::DataType& decoder_batch_ids_dtype,
     const paddle::DataType& decoder_tile_ids_per_batch_dtype,
     const paddle::DataType& decoder_num_blocks_dtype,
+    const paddle::DataType& decoder_num_blocks_cpu_dtype,
+    const paddle::DataType& decoder_chunk_size_cpu_dtype,
     const paddle::DataType& max_enc_len_this_time_dtype,
     const paddle::DataType& max_dec_len_this_time_dtype,
     const paddle::DataType& max_len_kv_dtype,
@@ -310,7 +412,7 @@ std::vector<paddle::DataType> MultiHeadLatentAttentionInferDtype(
     const float quant_max_bound,
     const float quant_min_bound,
     const float out_linear_in_scale,
-    const int speculate_max_draft_token_num,
+    const int speculate_draft_total_token_num,
     const bool causal,
     const bool speculate_decoder) {
   if (compute_dtype == "bf16") {
@@ -329,6 +431,7 @@ PD_BUILD_OP(multi_head_latent_attention)
              "seq_lens_encoder",
              "seq_lens_decoder",
              "seq_lens_this_time",
+             "cu_seqlens_q",
              "padding_offsets",
              "cum_offsets",
              "block_tables",
@@ -341,6 +444,8 @@ PD_BUILD_OP(multi_head_latent_attention)
              "decoder_batch_ids",
              "decoder_tile_ids_per_batch",
              "decoder_num_blocks",
+             "decoder_num_blocks_cpu",
+             "decoder_chunk_size_cpu",
              "max_enc_len_this_time",
              "max_dec_len_this_time",
              "max_len_kv",
@@ -364,7 +469,7 @@ PD_BUILD_OP(multi_head_latent_attention)
             "quant_max_bound: float",
             "quant_min_bound: float",
             "out_linear_in_scale: float",
-            "speculate_max_draft_token_num: int",
+            "speculate_draft_total_token_num: int",
             "causal: bool",
             "speculate_decoder: bool"})
     .SetKernelFn(PD_KERNEL(MultiHeadLatentAttention))

csrc/gpu/noaux_tc.cu

@@ -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.
+
+#pragma once
+
+#include <algorithm>
+#include <optional>
+
+#include "helper.h"
+#include "noauxtc_kernel.h"
+
+std::vector<paddle::Tensor> NoauxTc(paddle::Tensor& scores,
+                                    paddle::Tensor& scores_with_bias,
+                                    int n_group,
+                                    int topk_group,
+                                    int topk,
+                                    float routed_scaling_factor) {
+  auto input_shape = scores_with_bias.shape();
+  int64_t num_tokens = input_shape[0];
+  int64_t num_experts = input_shape[1];
+  auto input_type = scores_with_bias.dtype();
+  auto place = scores_with_bias.place();
+  auto group_scores = paddle::empty({num_tokens, n_group}, input_type, place);
+  auto stream = scores_with_bias.stream();
+
+  invokeNoAuxTc<float>(reinterpret_cast<float*>(scores.data<float>()),
+                       reinterpret_cast<float*>(group_scores.data<float>()),
+                       reinterpret_cast<float*>(scores_with_bias.data<float>()),
+                       num_tokens,
+                       num_experts,
+                       n_group,
+                       topk_group,
+                       topk,
+                       routed_scaling_factor,
+                       stream);
+
+  return {scores};
+}
+
+std::vector<paddle::DataType> NoauxTcInferDtype(
+    const paddle::DataType& scores_dtype,
+    const paddle::DataType& scores_with_bias_dtype) {
+  return {scores_dtype};
+}
+
+std::vector<std::vector<int64_t>> NoauxTcInferShape(
+    const std::vector<int64_t>& scores_shape,
+    const std::vector<int64_t>& gating_output_shape) {
+  return {scores_shape};
+}
+
+PD_BUILD_OP(noaux_tc)
+    .Inputs({"scores", "scores_with_bias"})
+    .Outputs({"output_tensor"})
+    .Attrs({"n_group: int",
+            "topk_group: int",
+            "topk:int",
+            "routed_scaling_factor: float"})
+    .SetKernelFn(PD_KERNEL(NoauxTc))
+    .SetInferShapeFn(PD_INFER_SHAPE(NoauxTcInferShape))
+    .SetInferDtypeFn(PD_INFER_DTYPE(NoauxTcInferDtype));

csrc/gpu/noauxtc_kernel.h

@@ -0,0 +1,551 @@
+// 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.
+
+// This code is partially inspired by and references the implementation found
+// in NVIDIA TRTLLM.
+
+#include <cooperative_groups.h>
+#include <cooperative_groups/reduce.h>
+
+namespace cg = cooperative_groups;
+
+constexpr unsigned FULL_WARP_MASK = 0xffffffff;
+constexpr int32_t WARP_SIZE = 32;
+constexpr int32_t BLOCK_SIZE = 512;
+constexpr int32_t NUM_WARPS_PER_BLOCK = BLOCK_SIZE / WARP_SIZE;
+
+namespace warp_topk {
+
+template <int size, typename T>
+__host__ __device__ constexpr T round_up_to_multiple_of(T len) {
+  if (len == 0) {
+    return 0;
+  }
+  return ((len - 1) / size + 1) * size;
+}
+
+template <typename T>
+constexpr __host__ __device__ bool isPowerOf2(T v) {
+  return (v && !(v & (v - 1)));
+}
+
+template <bool greater, typename T>
+__device__ bool is_better_than(T val, T baseline) {
+  return (val > baseline && greater) || (val < baseline && !greater);
+}
+
+template <typename T, typename idxT>
+int calc_smem_size_for_block_wide(int num_of_warp, int64_t k) {
+  int64_t cache_topk = (sizeof(T) + sizeof(idxT)) * num_of_warp * k;
+  int64_t n = std::max<int>(num_of_warp / 2 * k, num_of_warp * WARP_SIZE);
+  return max(cache_topk,
+             round_up_to_multiple_of<256>(n * sizeof(T)) + n * sizeof(idxT));
+}
+
+template <int size, bool ascending, typename T, typename idxT>
+struct BitonicMerge {
+  // input should be a bitonic sequence, and sort it to be a monotonic sequence
+  __device__ static void merge(T* __restrict__ val_arr,
+                               idxT* __restrict__ idx_arr) {
+    static_assert(isPowerOf2(size));
+    static_assert(size >= 2 * WARP_SIZE);
+    constexpr int arr_len = size / WARP_SIZE;
+
+    constexpr int stride = arr_len / 2;
+    for (int i = 0; i < stride; ++i) {
+      int const other_i = i + stride;
+      T& val = val_arr[i];
+      T& other_val = val_arr[other_i];
+      if ((val > other_val && ascending) || (val < other_val && !ascending)) {
+        T tmp = val;
+        val = other_val;
+        other_val = tmp;
+
+        idxT tmp2 = idx_arr[i];
+        idx_arr[i] = idx_arr[other_i];
+        idx_arr[other_i] = tmp2;
+      }
+    }
+
+    BitonicMerge<size / 2, ascending, T, idxT>::merge(val_arr, idx_arr);
+    BitonicMerge<size / 2, ascending, T, idxT>::merge(val_arr + arr_len / 2,
+                                                      idx_arr + arr_len / 2);
+  }
+};
+
+template <int size, bool ascending, typename T, typename idxT>
+struct BitonicSort {
+  __device__ static void sort(T* __restrict__ val_arr,
+                              idxT* __restrict__ idx_arr) {
+    static_assert(isPowerOf2(size));
+    static_assert(size >= 2 * WARP_SIZE);
+    constexpr int arr_len = size / WARP_SIZE;
+
+    BitonicSort<size / 2, true, T, idxT>::sort(val_arr, idx_arr);
+    BitonicSort<size / 2, false, T, idxT>::sort(val_arr + arr_len / 2,
+                                                idx_arr + arr_len / 2);
+    BitonicMerge<size, ascending, T, idxT>::merge(val_arr, idx_arr);
+  }
+};
+
+template <bool ascending, typename T, typename idxT>
+struct BitonicSort<32, ascending, T, idxT> {
+  __device__ static void sort(T* __restrict__ val_arr,
+                              idxT* __restrict__ idx_arr) {
+    int const lane = threadIdx.x % WARP_SIZE;
+
+    // ascending doesn't matter before merging since all we need is a bitonic
+    // sequence
+    for (int stage = 0; stage < 4; ++stage) {
+      for (int stride = (1 << stage); stride > 0; stride /= 2) {
+        bool reverse = (lane >> stage) & 2;
+        bool is_second = lane & stride;
+
+        T other = __shfl_xor_sync(FULL_WARP_MASK, *val_arr, stride);
+        idxT other_idx = __shfl_xor_sync(FULL_WARP_MASK, *idx_arr, stride);
+        if (*val_arr != other && (*val_arr > other) != (reverse != is_second)) {
+          *val_arr = other;
+          *idx_arr = other_idx;
+        }
+      }
+    }
+
+    BitonicMerge<32, ascending, T, idxT>::merge(val_arr, idx_arr);
+  }
+};
+
+template <bool ascending, typename T, typename idxT>
+struct BitonicMerge<32, ascending, T, idxT> {
+  __device__ static void merge(T* __restrict__ val_arr,
+                               idxT* __restrict__ idx_arr) {
+    int const lane = threadIdx.x % WARP_SIZE;
+    for (int stride = WARP_SIZE / 2; stride > 0; stride /= 2) {
+      bool is_second = lane & stride;
+      T& val = *val_arr;
+      T other = __shfl_xor_sync(FULL_WARP_MASK, val, stride);
+      idxT& idx = *idx_arr;
+      idxT other_idx = __shfl_xor_sync(FULL_WARP_MASK, idx, stride);
+      if (val != other && ((val > other) == (ascending != is_second))) {
+        val = other;
+        idx = other_idx;
+      }
+    }
+  }
+};
+
+template <int capacity, bool greater, typename T, typename idxT>
+class WarpSort {
+public:
+  __device__ WarpSort(idxT k, T dummy)
+      : lane_(threadIdx.x % WARP_SIZE), k_(k), dummy_(dummy) {
+    static_assert(capacity >= WARP_SIZE && isPowerOf2(capacity));
+
+    for (int i = 0; i < max_arr_len_; ++i) {
+      val_arr_[i] = dummy_;
+    }
+  }
+
+  // load and merge k sorted values
+  __device__ void load_sorted(T const* __restrict__ in,
+                              idxT const* __restrict__ in_idx,
+                              idxT start) {
+    idxT idx = start + WARP_SIZE - 1 - lane_;
+    for (int i = max_arr_len_ - 1; i >= 0; --i, idx += WARP_SIZE) {
+      if (idx < start + k_) {
+        T t = in[idx];
+        if (is_better_than<greater>(t, val_arr_[i])) {
+          val_arr_[i] = t;
+          idx_arr_[i] = in_idx[idx];
+        }
+      }
+    }
+
+    BitonicMerge<capacity, !greater, T, idxT>::merge(val_arr_, idx_arr_);
+  }
+
+  __device__ void dump(T* __restrict__ out, idxT* __restrict__ out_idx) const {
+    for (int i = 0; i < max_arr_len_; ++i) {
+      idxT out_i = i * WARP_SIZE + lane_;
+      if (out_i < k_) {
+        out[out_i] = val_arr_[i];
+        out_idx[out_i] = idx_arr_[i];
+      }
+    }
+  }
+
+  __device__ void dumpIdx(idxT* __restrict__ out_idx) const {
+    for (int i = 0; i < max_arr_len_; ++i) {
+      idxT out_i = i * WARP_SIZE + lane_;
+      if (out_i < k_) {
+        out_idx[out_i] = idx_arr_[i];
+      }
+    }
+  }
+
+protected:
+  static constexpr int max_arr_len_ = capacity / WARP_SIZE;
+
+  T val_arr_[max_arr_len_];
+  idxT idx_arr_[max_arr_len_];
+
+  int const lane_;
+  idxT const k_;
+  T const dummy_;
+
+};  // end class WarpSort
+
+template <int capacity, bool greater, typename T, typename idxT>
+class WarpSelect : public WarpSort<capacity, greater, T, idxT> {
+public:
+  __device__ WarpSelect(idxT k, T dummy)
+      : WarpSort<capacity, greater, T, idxT>(k, dummy),
+        k_th_(dummy),
+        k_th_lane_((k - 1) % WARP_SIZE) {
+    extern __shared__ char smem_buf[];  // extern __shared__ T smem_buf[];
+
+    int const num_of_warp = blockDim.x / WARP_SIZE;
+    int const warp_id = threadIdx.x / WARP_SIZE;
+    val_smem_ = reinterpret_cast<T*>(smem_buf);
+    val_smem_ += warp_id * WARP_SIZE;
+    idx_smem_ = reinterpret_cast<idxT*>(
+        smem_buf +
+        round_up_to_multiple_of<256>(num_of_warp * sizeof(T) * WARP_SIZE));
+    idx_smem_ += warp_id * WARP_SIZE;
+  }
+
+  __device__ void add(T const* in, idxT start, idxT end) {
+    idxT const end_for_fullwarp =
+        round_up_to_multiple_of<WARP_SIZE>(end - start) + start;
+    for (idxT i = start + lane_; i < end_for_fullwarp; i += WARP_SIZE) {
+      T val = (i < end) ? in[i] : dummy_;
+      add(val, i);
+    }
+  }
+
+  __device__ void add(T val, idxT idx) {
+    bool do_add = is_better_than<greater>(val, k_th_);
+    uint32_t mask = __ballot_sync(FULL_WARP_MASK, do_add);
+    if (mask == 0) {
+      return;
+    }
+
+    int pos = smem_buf_len_ + __popc(mask & ((0x1u << lane_) - 1));
+    if (do_add && pos < WARP_SIZE) {
+      val_smem_[pos] = val;
+      idx_smem_[pos] = idx;
+      do_add = false;
+    }
+    smem_buf_len_ += __popc(mask);
+    if (smem_buf_len_ >= WARP_SIZE) {
+      __syncwarp();
+      merge_buf_(val_smem_[lane_], idx_smem_[lane_]);
+      smem_buf_len_ -= WARP_SIZE;
+    }
+    if (do_add) {
+      pos -= WARP_SIZE;
+      val_smem_[pos] = val;
+      idx_smem_[pos] = idx;
+    }
+    __syncwarp();
+  }
+
+  __device__ void done() {
+    if (smem_buf_len_) {
+      T val = (lane_ < smem_buf_len_) ? val_smem_[lane_] : dummy_;
+      idxT idx = (lane_ < smem_buf_len_) ? idx_smem_[lane_] : 0;
+      merge_buf_(val, idx);
+    }
+
+    // after done(), smem is used for merging results among warps
+    __syncthreads();
+  }
+
+private:
+  __device__ void set_k_th_() {
+    k_th_ = __shfl_sync(FULL_WARP_MASK, val_arr_[max_arr_len_ - 1], k_th_lane_);
+  }
+
+  __device__ void merge_buf_(T val, idxT idx) {
+    BitonicSort<WARP_SIZE, greater, T, idxT>::sort(&val, &idx);
+
+    T& old = val_arr_[max_arr_len_ - 1];
+    if (is_better_than<greater>(val, old)) {
+      old = val;
+      idx_arr_[max_arr_len_ - 1] = idx;
+    }
+
+    BitonicMerge<capacity, !greater, T, idxT>::merge(val_arr_, idx_arr_);
+
+    set_k_th_();
+  }
+
+  using WarpSort<capacity, greater, T, idxT>::max_arr_len_;
+  using WarpSort<capacity, greater, T, idxT>::val_arr_;
+  using WarpSort<capacity, greater, T, idxT>::idx_arr_;
+  using WarpSort<capacity, greater, T, idxT>::lane_;
+  using WarpSort<capacity, greater, T, idxT>::k_;
+  using WarpSort<capacity, greater, T, idxT>::dummy_;
+
+  T* val_smem_;
+  idxT* idx_smem_;
+  int smem_buf_len_ = 0;
+
+  T k_th_;
+  int const k_th_lane_;
+};  // end class WarpSelect
+}  // namespace warp_topk
+
+template <typename T>
+__device__ void topk_with_k2(T* output,
+                             T const* input,
+                             cg::thread_block_tile<32> const& tile,
+                             int32_t const lane_id,
+                             int const num_experts_per_group) {
+  // Get the top2 per thread
+  T largest = cuda::std::numeric_limits<T>::min();
+  T second_largest = cuda::std::numeric_limits<T>::min();
+
+  if (num_experts_per_group > WARP_SIZE) {
+    for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
+      T value = input[i];
+      if (value > largest) {
+        second_largest = largest;
+        largest = value;
+      } else if (value > second_largest) {
+        second_largest = value;
+      }
+    }
+  } else {
+    for (int i = lane_id; i < num_experts_per_group; i += WARP_SIZE) {
+      largest = input[i];
+    }
+  }
+
+  __syncwarp();  // Ensure all threads have valid data before reduction
+  // Get the top2 warpwise
+  T max1 = cg::reduce(tile, largest, cg::greater<T>());
+
+  T max2 = max1;
+  bool equal_to_max1 = (max1 == largest);
+  int count_max1 = __popc(__ballot_sync(FULL_WARP_MASK, equal_to_max1));
+
+  if (count_max1 == 1) {
+    largest = (largest == max1) ? second_largest : largest;
+    max2 = cg::reduce(tile, largest, cg::greater<T>());
+  }
+
+  if (lane_id == 0) {
+    *output = max1 + max2;
+  }
+}
+
+template <typename T>
+__global__ void topk_with_k2_kernel(T* output,
+                                    T* input,
+                                    int64_t const num_tokens,
+                                    int64_t const num_cases,
+                                    int64_t const n_group,
+                                    int64_t const num_experts_per_group) {
+  int32_t warp_id = threadIdx.x / WARP_SIZE;
+  int32_t lane_id = threadIdx.x % WARP_SIZE;
+
+  int32_t case_id = blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id;
+  if (case_id < num_cases) {
+    input += case_id * num_experts_per_group;
+    output += case_id;
+
+    cg::thread_block block = cg::this_thread_block();
+    cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);
+
+    topk_with_k2(output, input, tile, lane_id, num_experts_per_group);
+  }
+}
+
+template <typename T>
+__global__ void group_idx_and_topk_idx_kernel(
+    T* scores,
+    T const* group_scores,
+    T* scores_with_bias,
+    int64_t const num_tokens,
+    int64_t const n_group,
+    int64_t const topk_group,
+    int64_t const topk,
+    int64_t const num_experts,
+    int64_t const num_experts_per_group,
+    double routed_scaling_factor) {
+  int32_t warp_id = threadIdx.x / WARP_SIZE;
+  int32_t lane_id = threadIdx.x % WARP_SIZE;
+  int32_t case_id =
+      blockIdx.x * NUM_WARPS_PER_BLOCK + warp_id;  // one per token
+  scores_with_bias += case_id * num_experts;
+  scores += case_id * num_experts;
+  group_scores += case_id * n_group;
+  int32_t align_num_experts_per_group =
+      warp_topk::round_up_to_multiple_of<WARP_SIZE>(num_experts_per_group);
+
+  cg::thread_block block = cg::this_thread_block();
+  cg::thread_block_tile<32> tile = cg::tiled_partition<32>(block);
+
+  extern __shared__ char smem_buf[];  // NOTE: reuse the shared memory here to
+                                      // store the target topk idx
+  int32_t* s_topk_idx = reinterpret_cast<int32_t*>(smem_buf) + warp_id * topk;
+  T* s_topk_value =
+      reinterpret_cast<T*>(s_topk_idx + NUM_WARPS_PER_BLOCK * topk) +
+      warp_id * topk;
+
+  T value = cuda::std::numeric_limits<T>::min();
+  T topk_group_value = cuda::std::numeric_limits<T>::min();
+  int32_t num_equalto_topkth_group;
+
+  if ((n_group > topk_group) && (case_id < num_tokens)) {
+    // calculate group_idx
+    int32_t target_num_min = WARP_SIZE - n_group + topk_group;
+    if (lane_id < n_group) {
+      value = group_scores[lane_id];
+    }
+
+    int count_equal_to_top_value = WARP_SIZE - n_group;
+    int pre_count_equal_to_top_value = 0;
+    // Use loop to find the largset top_group
+    while (count_equal_to_top_value < target_num_min) {
+      __syncwarp();  // Ensure all threads have valid data before reduction
+      topk_group_value = cg::reduce(tile, value, cg::greater<T>());
+      if (value == topk_group_value) {
+        value = cuda::std::numeric_limits<T>::min();
+      }
+      pre_count_equal_to_top_value = count_equal_to_top_value;
+      count_equal_to_top_value = __popc(__ballot_sync(
+          FULL_WARP_MASK, (value == cuda::std::numeric_limits<T>::min())));
+    }
+    num_equalto_topkth_group = target_num_min - pre_count_equal_to_top_value;
+  }
+  __syncthreads();
+
+  warp_topk::WarpSelect</*capability*/ WARP_SIZE, /*greater*/ true, T, int32_t>
+      queue((int32_t)topk, cuda::std::numeric_limits<T>::min());
+
+  int count_equalto_topkth_group = 0;
+  if (case_id < num_tokens) {
+    for (int i_group = 0; i_group < n_group; i_group++) {
+      if ((group_scores[i_group] > topk_group_value) ||
+          ((group_scores[i_group] == topk_group_value) &&
+           (count_equalto_topkth_group < num_equalto_topkth_group))) {
+        int32_t offset = i_group * num_experts_per_group;
+        for (int32_t i = lane_id; i < align_num_experts_per_group;
+             i += WARP_SIZE) {
+          T candidates = i < num_experts_per_group
+                             ? scores_with_bias[offset + i]
+                             : cuda::std::numeric_limits<T>::min();
+          queue.add(candidates, offset + i);
+        }
+        if (group_scores[i_group] == topk_group_value) {
+          count_equalto_topkth_group++;
+        }
+      }
+    }
+    queue.done();
+    __syncwarp();
+    // Get the topk_idx
+    queue.dumpIdx(s_topk_idx);
+    __syncwarp();
+  }
+
+  // Load the valid score value
+  // Calculate the summation
+  float topk_sum = 1e-20;
+  if (case_id < num_tokens) {
+    for (int i = lane_id;
+         i < warp_topk::round_up_to_multiple_of<WARP_SIZE>(topk);
+         i += WARP_SIZE) {
+      T value = i < topk ? scores[s_topk_idx[i]]
+                         : 0.0f;  // Load the valid value of expert
+      if (i < topk) {
+        s_topk_value[i] = value;
+      }
+      topk_sum += reduce(tile, value, cg::plus<float>());
+    }
+  }
+
+  __syncthreads();
+  if (case_id < num_tokens) {
+    for (int i = lane_id; i < num_experts; i += WARP_SIZE) {
+      scores[i] = 0;
+    }
+  }
+  __threadfence();
+  __syncthreads();
+
+  if (case_id < num_tokens) {
+    for (int i = lane_id; i < topk; i += WARP_SIZE) {
+      float value = s_topk_value[i] / topk_sum * routed_scaling_factor;
+      scores[s_topk_idx[i]] = value;
+    }
+  }
+}
+
+template <typename T>
+void invokeNoAuxTc(T* scores,
+                   T* group_scores,
+                   T* scores_with_bias,
+                   int64_t const num_tokens,
+                   int64_t const num_experts,
+                   int64_t const n_group,
+                   int64_t const topk_group,
+                   int64_t const topk,
+                   double const routed_scaling_factor,
+                   cudaStream_t const stream) {
+  int64_t num_cases = num_tokens * n_group;
+  int64_t topk_with_k2_num_blocks = (num_cases - 1) / NUM_WARPS_PER_BLOCK + 1;
+  topk_with_k2_kernel<T><<<topk_with_k2_num_blocks, BLOCK_SIZE, 0, stream>>>(
+      group_scores,
+      scores_with_bias,
+      num_tokens,
+      num_cases,
+      n_group,
+      num_experts / n_group);
+
+  int64_t topk_with_k_group_num_blocks =
+      (num_tokens - 1) / NUM_WARPS_PER_BLOCK + 1;
+  size_t dynamic_smem_in_bytes =
+      warp_topk::calc_smem_size_for_block_wide<T, int32_t>(NUM_WARPS_PER_BLOCK,
+                                                           topk);
+
+  group_idx_and_topk_idx_kernel<T><<<topk_with_k_group_num_blocks,
+                                     BLOCK_SIZE,
+                                     dynamic_smem_in_bytes,
+                                     stream>>>(scores,
+                                               group_scores,
+                                               scores_with_bias,
+                                               num_tokens,
+                                               n_group,
+                                               topk_group,
+                                               topk,
+                                               num_experts,
+                                               num_experts / n_group,
+                                               routed_scaling_factor);
+}
+
+#define INSTANTIATE_NOAUX_TC(T)                                      \
+  template void invokeNoAuxTc<T>(T * scores,                         \
+                                 T * group_scores,                   \
+                                 T * scores_with_bias,               \
+                                 int64_t const num_tokens,           \
+                                 int64_t const num_experts,          \
+                                 int64_t const n_group,              \
+                                 int64_t const topk_group,           \
+                                 int64_t const topk,                 \
+                                 double const routed_scaling_factor, \
+                                 cudaStream_t const stream);
+
+INSTANTIATE_NOAUX_TC(float);

csrc/gpu/quant_int8.cu

@@ -271,7 +271,7 @@ std::vector<paddle::DataType> QuantInt8Dtype(const paddle::DataType& input_dtype
 }
 
 PD_BUILD_OP(quant_int8)
-    .Inputs({"intput", paddle::Optional("shift"),paddle::Optional("smooth") })
+    .Inputs({"input", paddle::Optional("shift"),paddle::Optional("smooth") })
     .Outputs({"output"})
     .Attrs({"scale: float","round_type: int","max_bound: float", "min_bound: float"})
     .SetKernelFn(PD_KERNEL(QuantInt8))

csrc/gpu/sage_attn_kernels/sageattn_fused.cuh

@@ -0,0 +1,19 @@
+#pragma once
+#include "paddle/extension.h"
+
+std::vector<paddle::Tensor> per_warp_int8_cuda(paddle::Tensor& q,
+                                            paddle::Tensor& k,
+                                            paddle::Tensor& km,
+                                            int BLKQ,
+                                            int WARPQ,
+                                            int BLKK,
+                                            int tensor_layout);
+                    
+std::vector<paddle::Tensor> per_channel_fp8(paddle::Tensor& v,
+                                            int tensor_layout,
+                                            float scale_max,
+                                            bool smooth_v);
+
+std::vector<paddle::Tensor> sub_mean(paddle::Tensor& v,
+                                    paddle::Tensor& vm,
+                                    int tensor_layout);

csrc/gpu/set_preids_token_penalty_multi_scores.cu

@@ -46,13 +46,13 @@ __global__ void set_preids_token_penalty_multi_scores_kernel(const bool *stop_fl
         const int64_t *input_ids_now = input_ids + tid * length_input_ids;
         const int seq_len_dec = seq_lens_decoder[tid];
         const int seq_len_enc = seq_lens_encoder[tid];
-        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stoped
+        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stopped
 
         const int step_idx_now = step_idx[bi];
         if (tid == 0 && step_idx_now >= 0) {
             if (seq_len_enc > 0) { // encoder, get last token accord to seq_lens_encoder
                 pre_ids_now[step_idx_now] = input_ids_now[seq_len_enc - 1];
-            } else { // decoedr, get first token
+            } else { // decoder, get first token
                 pre_ids_now[step_idx_now] = input_ids_now[0];
             }
         }

csrc/gpu/set_value_by_flags_v2.cu

@@ -29,11 +29,11 @@ __global__ void set_value_by_flag_and_id_v2(const bool *stop_flags,
         const int64_t *input_ids_now = input_ids + tid * length_input_ids;
         const int seq_len_dec = seq_lens_decoder[tid];
         const int seq_len_enc = seq_lens_encoder[tid];
-        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stoped
+        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stopped
         if (step_idx[tid] >= 0) {
             if (seq_len_dec == 0) { // encoder, get last token accord to seq_lens_encoder
                 pre_ids_all_now[step_idx[tid]] = input_ids_now[seq_len_enc - 1];
-            } else { // decoedr, get first token
+            } else { // decoder, get first token
                 pre_ids_all_now[step_idx[tid]] = input_ids_now[0];
             }
         }

csrc/gpu/speculate_decoding_kernels/speculate_set_value_by_flags.cu

@@ -30,7 +30,7 @@ __global__ void speculate_set_value_by_flag_and_id(int64_t *pre_ids_all,
         const int64_t *accept_tokens_now = accept_tokens + tid * max_draft_tokens;
         const int seq_len_dec = seq_lens_decoder[tid];
         const int seq_len_enc = seq_lens_encoder[tid];
-        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stoped
+        if (seq_len_dec == 0 && seq_len_enc == 0) return; // stopped
         if (step_idx[tid] >= 0) {
             for (int i = 0; i < accept_num[tid]; i++) {
                 pre_ids_all_now[step_idx[tid] - i] = accept_tokens_now[accept_num[tid] - 1 - i];

csrc/gpu/step.cu

@@ -1,11 +1,11 @@
 // 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.
@@ -16,13 +16,27 @@
 
 // #define DEBUG_STEP
 
+__device__ bool in_need_block_list(const int &qid,
+                                   int *need_block_list,
+                                   const int &need_block_len) {
+  bool res = false;
+  for (int i = 0; i < need_block_len; i++) {
+    if (qid == need_block_list[i]) {
+      res = true;
+      need_block_list[i] = -1;
+      break;
+    }
+  }
+  return res;
+}
+
 __global__ void free_and_dispatch_block(bool *stop_flags,
                                         int *seq_lens_this_time,
                                         int *seq_lens_decoder,
                                         int *block_tables,
                                         int *encoder_block_lens,
                                         bool *is_block_step,
-                                        int *step_block_list, // [bsz]
+                                        int *step_block_list,  // [bsz]
                                         int *step_len,
                                         int *recover_block_list,
                                         int *recover_len,
@@ -35,124 +49,183 @@ __global__ void free_and_dispatch_block(bool *stop_flags,
                                         const int bsz,
                                         const int block_size,
                                         const int block_num_per_seq,
-                                        const int max_decoder_block_num,
-                                        const int speculate_step_token_num) {
-    typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 512> BlockReduce;
-    __shared__ typename BlockReduce::TempStorage temp_storage;
-    const int tid = threadIdx.x;
-    if (tid < bsz) {
-        int *block_table_now = block_tables + tid * block_num_per_seq;
-        if (stop_flags[tid] && !is_block_step[tid]) {
-            // 回收block块
-            first_token_ids[tid] = -1;
-            const int encoder_block_len = encoder_block_lens[tid];
-            const int decoder_used_len = used_list_len[tid];
-            if (decoder_used_len > 0) {
-                const int ori_free_list_len = atomicAdd(free_list_len, decoder_used_len);
-#ifdef DEBUG_STEP
-                printf("free block seq_id: %d, free block num: %d, encoder_block_len: %d, ori_free_list_len: %d\n", 
-                        tid, decoder_used_len, encoder_block_len, ori_free_list_len);
-#endif
-                for (int i = 0; i < decoder_used_len; i++) {
-                    free_list[ori_free_list_len + i] = block_table_now[encoder_block_len + i];
-                    block_table_now[encoder_block_len + i] = -1;
-                }
-                encoder_block_lens[tid] = 0;
-                used_list_len[tid] = 0;
-            }
-        } else if (seq_lens_decoder[tid] != 0 && block_table_now[(seq_lens_decoder[tid] + speculate_step_token_num) / block_size] == -1) {
-            // 统计需要分配block的位置和总数
-            const int ori_need_block_len = atomicAdd(need_block_len, 1);
-            need_block_list[ori_need_block_len] = tid;
-#ifdef DEBUG_STEP
-            printf("seq_id: %d need block\n", tid);
-#endif
-        }
-    }
-    __syncthreads();
+                                        const int max_decoder_block_num) {
+  typedef cub::BlockReduce<cub::KeyValuePair<int, int>, 256> BlockReduce;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ bool step_max_block_flag;
+  __shared__ int in_need_block_list_len;
+  const int tid = threadIdx.x;
+  if (tid < bsz) {
     if (tid == 0) {
-        printf("need_block_len: %d, free_list_len: %d\n", need_block_len[0], free_list_len[0]);
+      step_max_block_flag = false;
+      in_need_block_list_len = 0;
     }
+    int *block_table_now = block_tables + tid * block_num_per_seq;
+    if (stop_flags[tid] && !is_block_step[tid]) {
+      // 回收block块
+      first_token_ids[tid] = -1;
+      const int encoder_block_len = encoder_block_lens[tid];
+      const int decoder_used_len = used_list_len[tid];
+      if (decoder_used_len > 0) {
+        const int ori_free_list_len =
+            atomicAdd(free_list_len, decoder_used_len);
+#ifdef DEBUG_STEP
+        printf(
+            "free block seq_id: %d, free block num: %d, "
+            "encoder_block_len: %d, ori_free_list_len: %d\n",
+            tid,
+            decoder_used_len,
+            encoder_block_len,
+            ori_free_list_len);
+#endif
+        for (int i = 0; i < decoder_used_len; i++) {
+          free_list[ori_free_list_len + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
+        }
+        encoder_block_lens[tid] = 0;
+        used_list_len[tid] = 0;
+      }
+    } else if (seq_lens_this_time[tid] != 0 &&
+               block_table_now[seq_lens_decoder[tid] / block_size] == -1) {
+      // 统计需要分配block的位置和总数
+#ifdef DEBUG_STEP
+      printf("step seq_id:%d, ##### pin 1 #####\n", tid);
+#endif
+      const int ori_need_block_len = atomicAdd(need_block_len, 1);
+      need_block_list[ori_need_block_len] = tid;
+#ifdef DEBUG_STEP
+      printf("seq_id: %d need block\n", tid);
+#endif
+    }
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 2 #####\n", tid);
+#endif
+  __syncthreads();
 
-    while (need_block_len[0] > free_list_len[0]) {
-        
-#ifdef DEBUG_STEP
-        if (tid == 0) {
-            printf("need_block_len: %d, free_list_len: %d\n", need_block_len[0], free_list_len[0]);
-        }
-#endif
-        // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len
-        const int used_block_num = tid < bsz && !is_block_step[tid]? used_list_len[tid] : 0;
-        cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
-        kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
-        if (tid == 0) {
-            const int encoder_block_len = encoder_block_lens[kv_pair.key];
-#ifdef DEBUG_STEP
-            printf("max_id: %d, max_num: %d, encoder_block_len: %d\n", kv_pair.key, kv_pair.value, encoder_block_len);
-#endif
-            int *block_table_now = block_tables + kv_pair.key * block_num_per_seq;
-            for (int i = 0; i < kv_pair.value; i++) {
-                free_list[free_list_len[0] + i] = block_table_now[encoder_block_len + i];
-                block_table_now[encoder_block_len + i] = -1;
-            }
-            step_block_list[step_len[0]] = kv_pair.key;
-            step_len[0] += 1;
-            free_list_len[0] += kv_pair.value;
-            stop_flags[kv_pair.key] = true;
-            is_block_step[kv_pair.key] = true;
-            seq_lens_this_time[kv_pair.key] = 0;
-            seq_lens_decoder[kv_pair.key] = 0;
-        }
-        __syncthreads();
+  while (need_block_len[0] > free_list_len[0]) {
+    if (tid == 0) {
+      printf("need_block_len: %d, free_list_len: %d\n",
+             need_block_len[0],
+             free_list_len[0]);
     }
-    // 为需要block的位置分配block，每个位置分配一个block
-    if (tid < need_block_len[0]) {
-        const int need_block_id = need_block_list[tid];
-        if (!stop_flags[need_block_id]) {
-            // 如果需要的位置正好是上一步中被释放的位置，不做处理
-            used_list_len[need_block_id] += 1;
-            const int ori_free_list_len = atomicSub(free_list_len, 1);
-            int *block_table_now = block_tables + need_block_id * block_num_per_seq;
-            block_table_now[(seq_lens_decoder[need_block_id] + speculate_step_token_num) / block_size] = free_list[ori_free_list_len - 1];
+    // 调度block，根据used_list_len从大到小回收block，直到满足need_block_len，已解码到最后一个block的query不参与调度（马上就结束）
+    const int used_block_num =
+        tid < bsz && !is_block_step[tid] &&
+                (step_max_block_flag ||
+                 used_list_len[tid] != max_decoder_block_num)
+            ? used_list_len[tid]
+            : 0;
+    cub::KeyValuePair<int, int> kv_pair = {tid, used_block_num};
+    kv_pair = BlockReduce(temp_storage).Reduce(kv_pair, cub::ArgMax());
+
+    if (tid == 0) {
+      if (kv_pair.value == 0) {
+        step_max_block_flag = true;
+      } else {
+        const int encoder_block_len = encoder_block_lens[kv_pair.key];
+        printf("step max_id: %d, max_num: %d, encoder_block_len: %d\n",
+               kv_pair.key,
+               kv_pair.value,
+               encoder_block_len);
+        int *block_table_now = block_tables + kv_pair.key * block_num_per_seq;
+        for (int i = 0; i < kv_pair.value; i++) {
+          free_list[free_list_len[0] + i] =
+              block_table_now[encoder_block_len + i];
+          block_table_now[encoder_block_len + i] = -1;
         }
-        need_block_list[tid] = -1;
+        step_block_list[step_len[0]] = kv_pair.key;
+        if (in_need_block_list(kv_pair.key,
+                               need_block_list,
+                               need_block_len[0] + in_need_block_list_len)) {
+          need_block_len[0] -= 1;
+          in_need_block_list_len += 1;
+        }
+        step_len[0] += 1;
+        free_list_len[0] += kv_pair.value;
+        stop_flags[kv_pair.key] = true;
+        is_block_step[kv_pair.key] = true;
+        seq_lens_this_time[kv_pair.key] = 0;
+        seq_lens_decoder[kv_pair.key] = 0;
+      }
     }
     __syncthreads();
-    // 计算可以复原的query id
-    if (tid == 0) {
-        int ori_free_list_len = free_list_len[0];
-        int ori_step_len = step_len[0];
-        if (ori_step_len > 0) {
-            int ori_step_block_id = step_block_list[ori_step_len - 1];
-            int tmp_used_len = used_list_len[ori_step_block_id];
-            // 比之前调度时多分配一个block，防止马上恢复刚调度的query(比如回收的seq_id在need_block_list中）
-            int used_len = tmp_used_len < max_decoder_block_num ? tmp_used_len + 1 : tmp_used_len;
-            while (ori_step_len > 0 && ori_free_list_len >= used_len) {
+  }
 #ifdef DEBUG_STEP
-                printf("recover seq_id: %d, free_list_len: %d, used_list_len: %d\n", 
-                        ori_step_block_id, ori_free_list_len, used_len);
+  printf("step seq_id:%d, ##### pin 3 #####\n", tid);
 #endif
-                recover_block_list[recover_len[0]] = ori_step_block_id;
-                is_block_step[ori_step_block_id] = false;
-                used_list_len[ori_step_block_id] = used_len;
-                ori_free_list_len -= used_len;
-                step_block_list[ori_step_len - 1] = -1;
-                step_len[0] -= 1;
-                recover_len[0] += 1;
-                ori_step_len = step_len[0];
-                if (ori_step_len > 0) {
-                    ori_step_block_id = step_block_list[ori_step_len - 1];
-                    tmp_used_len = used_list_len[ori_step_block_id];
-                    used_len = tmp_used_len < max_decoder_block_num ? tmp_used_len + 1 : tmp_used_len;
-                }
-            }
-        }
-        need_block_len[0] = 0;
+  // 为需要block的位置分配block，每个位置分配一个block
+  if (tid < need_block_len[0] + in_need_block_list_len) {
+    const int need_block_id = need_block_list[tid];
+    if (need_block_id != -1) {
+      if (!stop_flags[need_block_id]) {
+        // 如果需要的位置正好是上一步中被释放的位置，不做处理
+        used_list_len[need_block_id] += 1;
+        const int ori_free_list_len = atomicSub(free_list_len, 1);
+        int *block_table_now = block_tables + need_block_id * block_num_per_seq;
+        block_table_now[seq_lens_decoder[need_block_id] / block_size] =
+            free_list[ori_free_list_len - 1];
+      }
+      need_block_list[tid] = -1;
     }
+  }
+  __syncthreads();
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 4 #####\n", tid);
+#endif
+  // 计算可以复原的query id
+  if (tid == 0) {
+    int ori_free_list_len = free_list_len[0];
+    int ori_step_len = step_len[0];
+    if (ori_step_len > 0) {
+      int ori_step_block_id = step_block_list[ori_step_len - 1];
+      int tmp_used_len = used_list_len[ori_step_block_id];
+      // 比之前调度时多分配一个block，防止马上恢复刚调度的query(比如回收的seq_id在need_block_list中）
+      const int max_decoder_block_num_this_seq =
+          max_decoder_block_num - encoder_block_lens[ori_step_block_id];
+      int used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
+                         ? tmp_used_len + 1
+                         : max_decoder_block_num_this_seq;
+#ifdef DEBUG_STEP
+      printf("#### ori_step_len:%d, ori_free_list_len:%d, used_len:%d  \n",
+             ori_step_len,
+             ori_free_list_len,
+             used_len);
+#endif
+      if (ori_step_len > 0 && ori_free_list_len >= used_len) {
+        printf(
+            "recover seq_id: %d, free_list_len: %d, "
+            "used_list_len: %d\n",
+            ori_step_block_id,
+            ori_free_list_len,
+            used_len);
+        recover_block_list[recover_len[0]] = ori_step_block_id;
+        is_block_step[ori_step_block_id] = false;
+        used_list_len[ori_step_block_id] = used_len;
+        ori_free_list_len -= used_len;
+        step_block_list[ori_step_len - 1] = -1;
+        step_len[0] -= 1;
+        recover_len[0] += 1;
+        ori_step_len = step_len[0];
+        if (ori_step_len > 0) {
+          ori_step_block_id = step_block_list[ori_step_len - 1];
+          tmp_used_len = used_list_len[ori_step_block_id];
+          used_len = tmp_used_len + 1 < max_decoder_block_num_this_seq
+                         ? tmp_used_len + 1
+                         : max_decoder_block_num_this_seq;
+        }
+      }
+    }
+    need_block_len[0] = 0;
+  }
+#ifdef DEBUG_STEP
+  printf("step seq_id:%d, ##### pin 5 #####\n", tid);
+#endif
 }
 
 // 根据上一步计算出的可以复原的query_id进行状态恢复
-__global__ void recover_block(int *recover_block_list, // [bsz]
+__global__ void recover_block(int *recover_block_list,  // [bsz]
                               int *recover_len,
                               bool *stop_flags,
                               int *seq_lens_this_time,
@@ -173,107 +246,157 @@ __global__ void recover_block(int *recover_block_list, // [bsz]
                               const int block_num_per_seq,
                               const int length,
                               const int pre_id_length) {
-    const int bid = blockIdx.x;
-    const int tid = threadIdx.x;
-    __shared__ int ori_free_list_len;
-    if (bid < recover_len[0]) {
-        const int recover_id = recover_block_list[bid];
-        const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
-        const int step_idx_now = step_idx[recover_id];
-        const int seq_len = ori_seq_len_encoder + step_idx_now;
-        const int encoder_block_len = encoder_block_lens[recover_id];
-        const int decoder_used_len = used_list_len[recover_id];
-        int *block_table_now = block_tables + recover_id * block_num_per_seq;
-        int64_t *input_ids_now = input_ids + recover_id * length;
-        int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
-        if (tid == 0) {
-            seq_lens_this_time[recover_id] = seq_len;
-            seq_lens_encoder[recover_id] = seq_len;
-            stop_flags[recover_id] = false;
-            input_ids_now[ori_seq_len_encoder + step_idx_now - 1] = next_tokens[recover_id]; // next tokens
-            input_ids_now[0] =
-                first_token_ids[recover_id];  // set first prompt token
-            const int ori_free_list_len_tid0 = atomicSub(free_list_len, decoder_used_len);
-            ori_free_list_len = ori_free_list_len_tid0;
+  const int bid = blockIdx.x;
+  const int tid = threadIdx.x;
+  __shared__ int ori_free_list_len;
+  if (bid < recover_len[0]) {
+    const int recover_id = recover_block_list[bid];
+    const int ori_seq_len_encoder = ori_seq_lens_encoder[recover_id];
+    const int step_idx_now = step_idx[recover_id];
+    const int seq_len = ori_seq_len_encoder + step_idx_now;
+    const int encoder_block_len = encoder_block_lens[recover_id];
+    const int decoder_used_len = used_list_len[recover_id];
+    int *block_table_now = block_tables + recover_id * block_num_per_seq;
+    int64_t *input_ids_now = input_ids + recover_id * length;
+    int64_t *pre_ids_now = pre_ids + recover_id * pre_id_length;
+    if (tid == 0) {
+      seq_lens_this_time[recover_id] = seq_len;
+      seq_lens_encoder[recover_id] = seq_len;
+      stop_flags[recover_id] = false;
+      input_ids_now[ori_seq_len_encoder + step_idx_now - 1] =
+          next_tokens[recover_id];                     // next tokens
+      input_ids_now[0] = first_token_ids[recover_id];  // set first prompt token
+      const int ori_free_list_len_tid0 =
+          atomicSub(free_list_len, decoder_used_len);
+      ori_free_list_len = ori_free_list_len_tid0;
 #ifdef DEBUG_STEP
-            printf("seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, seq_len: %d, ori_free_list_len_tid0: %d, ori_free_list_len: %d\n", 
-                    recover_id, ori_seq_len_encoder, step_idx_now, seq_len, ori_free_list_len_tid0, ori_free_list_len);
+      printf(
+          "seq_id: %d, ori_seq_len_encoder: %d, step_idx_now: %d, "
+          "seq_len: %d, ori_free_list_len_tid0: %d, "
+          "ori_free_list_len: %d\n",
+          recover_id,
+          ori_seq_len_encoder,
+          step_idx_now,
+          seq_len,
+          ori_free_list_len_tid0,
+          ori_free_list_len);
 #endif
-        }
-        __syncthreads();
-        // 恢复block table
-        for (int i = tid; i < decoder_used_len; i += blockDim.x) {
-            block_table_now[encoder_block_len + i] = free_list[ori_free_list_len - i - 1];
-        }
-        // 恢复input_ids
-        for (int i = tid; i < step_idx_now - 1; i += blockDim.x) {
-            input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
-        }
     }
+    __syncthreads();
+    // 恢复block table
+    for (int i = tid; i < decoder_used_len; i += blockDim.x) {
+      block_table_now[encoder_block_len + i] =
+          free_list[ori_free_list_len - i - 1];
+    }
+    // 恢复input_ids
+    for (int i = tid; i < step_idx_now - 1; i += blockDim.x) {
+      input_ids_now[ori_seq_len_encoder + i] = pre_ids_now[i + 1];
+    }
+  }
 
-    if (bid == 0 && tid == 0) {
-        recover_len[0] = 0;
-    }
+  if (bid == 0 && tid == 0) {
+    recover_len[0] = 0;
+  }
 }
 
-void StepPaddle(const paddle::Tensor& stop_flags,
-                const paddle::Tensor& seq_lens_this_time,
-                const paddle::Tensor& ori_seq_lens_encoder,
-                const paddle::Tensor& seq_lens_encoder,
-                const paddle::Tensor& seq_lens_decoder,
-                const paddle::Tensor& block_tables, // [bsz, block_num_per_seq]
-                const paddle::Tensor& encoder_block_lens,
-                const paddle::Tensor& is_block_step,
-                const paddle::Tensor& step_block_list,
-                const paddle::Tensor& step_lens,
-                const paddle::Tensor& recover_block_list,
-                const paddle::Tensor& recover_lens,
-                const paddle::Tensor& need_block_list,
-                const paddle::Tensor& need_block_len,
-                const paddle::Tensor& used_list_len, 
-                const paddle::Tensor& free_list,
-                const paddle::Tensor& free_list_len,
-                const paddle::Tensor& input_ids,
-                const paddle::Tensor& pre_ids,
-                const paddle::Tensor& step_idx,
-                const paddle::Tensor& next_tokens,
+void StepPaddle(const paddle::Tensor &stop_flags,
+                const paddle::Tensor &seq_lens_this_time,
+                const paddle::Tensor &ori_seq_lens_encoder,
+                const paddle::Tensor &seq_lens_encoder,
+                const paddle::Tensor &seq_lens_decoder,
+                const paddle::Tensor &block_tables,  // [bsz, block_num_per_seq]
+                const paddle::Tensor &encoder_block_lens,
+                const paddle::Tensor &is_block_step,
+                const paddle::Tensor &step_block_list,
+                const paddle::Tensor &step_lens,
+                const paddle::Tensor &recover_block_list,
+                const paddle::Tensor &recover_lens,
+                const paddle::Tensor &need_block_list,
+                const paddle::Tensor &need_block_len,
+                const paddle::Tensor &used_list_len,
+                const paddle::Tensor &free_list,
+                const paddle::Tensor &free_list_len,
+                const paddle::Tensor &input_ids,
+                const paddle::Tensor &pre_ids,
+                const paddle::Tensor &step_idx,
+                const paddle::Tensor &next_tokens,
                 const paddle::Tensor &first_token_ids,
                 const int block_size,
-                const int encoder_decoder_block_num,
-                const int speculate_step_token_num) {
-    auto cu_stream = seq_lens_this_time.stream();
-    const int bsz = seq_lens_this_time.shape()[0];
-    const int block_num_per_seq = block_tables.shape()[1];
-    const int length = input_ids.shape()[1];
-    const int pre_id_length = pre_ids.shape()[1];
-    constexpr int BlockSize = 512; // bsz < 256
-    const int max_decoder_block_num = length / block_size;
+                const int encoder_decoder_block_num) {
+  auto cu_stream = seq_lens_this_time.stream();
+  const int bsz = seq_lens_this_time.shape()[0];
+  const int block_num_per_seq = block_tables.shape()[1];
+  const int length = input_ids.shape()[1];
+  const int pre_id_length = pre_ids.shape()[1];
+  constexpr int BlockSize = 256;  // bsz <= 256
+  const int max_decoder_block_num = length / block_size;
+  // const int max_decoder_block_num = 2048 / block_size -
+  // encoder_decoder_block_num;
 #ifdef DEBUG_STEP
-    printf("bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: %d\n", bsz, block_num_per_seq, length, max_decoder_block_num);
+  printf(
+      "bsz: %d, block_num_per_seq: %d, length: %d, max_decoder_block_num: "
+      "%d\n",
+      bsz,
+      block_num_per_seq,
+      length,
+      max_decoder_block_num);
 #endif
-    free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
-        const_cast<bool*>(stop_flags.data<bool>()),
-        const_cast<int*>(seq_lens_this_time.data<int>()),
-        const_cast<int*>(seq_lens_decoder.data<int>()),
-        const_cast<int*>(block_tables.data<int>()),
-        const_cast<int*>(encoder_block_lens.data<int>()),
-        const_cast<bool*>(is_block_step.data<bool>()),
-        const_cast<int*>(step_block_list.data<int>()),
-        const_cast<int*>(step_lens.data<int>()),
-        const_cast<int*>(recover_block_list.data<int>()),
-        const_cast<int*>(recover_lens.data<int>()),
-        const_cast<int*>(need_block_list.data<int>()),
-        const_cast<int*>(need_block_len.data<int>()),
-        const_cast<int*>(used_list_len.data<int>()),
-        const_cast<int*>(free_list.data<int>()),
-        const_cast<int*>(free_list_len.data<int>()),
-        const_cast<int64_t *>(first_token_ids.data<int64_t>()),
+  free_and_dispatch_block<<<1, BlockSize, 0, cu_stream>>>(
+      const_cast<bool *>(stop_flags.data<bool>()),
+      const_cast<int *>(seq_lens_this_time.data<int>()),
+      const_cast<int *>(seq_lens_decoder.data<int>()),
+      const_cast<int *>(block_tables.data<int>()),
+      const_cast<int *>(encoder_block_lens.data<int>()),
+      const_cast<bool *>(is_block_step.data<bool>()),
+      const_cast<int *>(step_block_list.data<int>()),
+      const_cast<int *>(step_lens.data<int>()),
+      const_cast<int *>(recover_block_list.data<int>()),
+      const_cast<int *>(recover_lens.data<int>()),
+      const_cast<int *>(need_block_list.data<int>()),
+      const_cast<int *>(need_block_len.data<int>()),
+      const_cast<int *>(used_list_len.data<int>()),
+      const_cast<int *>(free_list.data<int>()),
+      const_cast<int *>(free_list_len.data<int>()),
+      const_cast<int64_t *>(first_token_ids.data<int64_t>()),
+      bsz,
+      block_size,
+      block_num_per_seq,
+      max_decoder_block_num);
+#ifdef DEBUG_STEP
+#ifdef PADDLE_WITH_HIP
+  hipDeviceSynchronize();
+#else
+  cudaDeviceSynchronize();
+#endif
+#endif
+  auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
+  const int grid_size = cpu_recover_lens.data<int>()[0];
+#ifdef DEBUG_STEP
+  printf("grid_size2 %d\n", grid_size);
+#endif
+  if (grid_size > 0) {
+    recover_block<<<grid_size, BlockSize, 0, cu_stream>>>(
+        const_cast<int *>(recover_block_list.data<int>()),
+        const_cast<int *>(recover_lens.data<int>()),
+        const_cast<bool *>(stop_flags.data<bool>()),
+        const_cast<int *>(seq_lens_this_time.data<int>()),
+        const_cast<int *>(ori_seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_encoder.data<int>()),
+        const_cast<int *>(seq_lens_decoder.data<int>()),
+        const_cast<int *>(block_tables.data<int>()),
+        const_cast<int *>(free_list.data<int>()),
+        const_cast<int *>(free_list_len.data<int>()),
+        const_cast<int64_t *>(input_ids.data<int64_t>()),
+        const_cast<int64_t *>(pre_ids.data<int64_t>()),
+        const_cast<int64_t *>(step_idx.data<int64_t>()),
+        const_cast<int *>(encoder_block_lens.data<int>()),
+        const_cast<int *>(used_list_len.data<int>()),
+        next_tokens.data<int64_t>(),
+        first_token_ids.data<int64_t>(),
         bsz,
-        block_size,
         block_num_per_seq,
-        max_decoder_block_num,
-        speculate_step_token_num
-    );
+        length,
+        pre_id_length);
 #ifdef DEBUG_STEP
 #ifdef PADDLE_WITH_HIP
     hipDeviceSynchronize();
@@ -281,47 +404,11 @@ void StepPaddle(const paddle::Tensor& stop_flags,
     cudaDeviceSynchronize();
 #endif
 #endif
-    auto cpu_recover_lens = recover_lens.copy_to(paddle::CPUPlace(), false);
-    const int grid_size = cpu_recover_lens.data<int>()[0];
-#ifdef DEBUG_STEP
-    printf("grid_size2 %d\n", grid_size);
-#endif
-    if (grid_size > 0) {
-        recover_block<<<grid_size, BlockSize, 0, cu_stream>>>(
-            const_cast<int*>(recover_block_list.data<int>()),
-            const_cast<int*>(recover_lens.data<int>()),
-            const_cast<bool*>(stop_flags.data<bool>()),
-            const_cast<int*>(seq_lens_this_time.data<int>()),
-            const_cast<int*>(ori_seq_lens_encoder.data<int>()),
-            const_cast<int*>(seq_lens_encoder.data<int>()),
-            const_cast<int*>(seq_lens_decoder.data<int>()),
-            const_cast<int*>(block_tables.data<int>()),
-            const_cast<int*>(free_list.data<int>()),
-            const_cast<int*>(free_list_len.data<int>()),
-            const_cast<int64_t*>(input_ids.data<int64_t>()),
-            const_cast<int64_t*>(pre_ids.data<int64_t>()),
-            const_cast<int64_t*>(step_idx.data<int64_t>()),
-            const_cast<int*>(encoder_block_lens.data<int>()),
-            const_cast<int*>(used_list_len.data<int>()),
-            next_tokens.data<int64_t>(),
-            first_token_ids.data<int64_t>(),
-            bsz,
-            block_num_per_seq,
-            length,
-            pre_id_length
-        );
-#ifdef DEBUG_STEP
-#ifdef PADDLE_WITH_HIP
-        hipDeviceSynchronize();
-#else
-        cudaDeviceSynchronize();
-#endif
-#endif
-    }
+  }
 }
 
 PD_BUILD_OP(step_paddle)
-    .Inputs({"stop_flags", 
+    .Inputs({"stop_flags",
              "seq_lens_this_time",
              "ori_seq_lens_encoder",
              "seq_lens_encoder",
@@ -342,10 +429,8 @@ PD_BUILD_OP(step_paddle)
              "pre_ids",
              "step_idx",
              "next_tokens",
-             "first_token_ids",})
-    .Attrs({"block_size: int",
-            "encoder_decoder_block_num: int",
-            "speculate_step_token_num: int"})
+             "first_token_ids"})
+    .Attrs({"block_size: int", "encoder_decoder_block_num: int"})
     .Outputs({"stop_flags_out",
               "seq_lens_this_time_out",
               "seq_lens_encoder_out",
@@ -363,7 +448,7 @@ PD_BUILD_OP(step_paddle)
               "free_list_out",
               "free_list_len_out",
               "input_ids_out",
-              "first_token_ids_out",})
+              "first_token_ids_out"})
     .SetInplaceMap({{"stop_flags", "stop_flags_out"},
                     {"seq_lens_this_time", "seq_lens_this_time_out"},
                     {"seq_lens_encoder", "seq_lens_encoder_out"},

csrc/gpu/unittest/run_mla_benchmark.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.
+
+#请使用cuda-12.8编译和运行benchnark
+export FLAGS_mla_use_tensorcore=1
+export FLAGS_mla_use_wg4=0
+export FLAGS_mla_dec_chunk_size=-1
+
+# python test_absorb_mla.py
+
+ncu=`which ncu`
+${ncu} --section regex:'^(?!Nvlink)' -o profile_mla_kernel -f --import-source on --cache-control=all --clock-control=base -k regex:"MLAWithKVCacheKernel|merge_multi_chunks_kernel" --print-source=cuda,sass --page source python test_absorb_mla.py

csrc/gpu/unittest/test_absorb_mla.py

@@ -0,0 +1,301 @@
+# 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 time
+
+import numpy as np
+import paddle
+import paddlenlp_ops
+
+np.random.seed(2024)
+paddle.seed(2024)
+
+
+def div_up(a, b):
+    return (a + b - 1) // b
+
+
+def alloc_diff(x, y):
+    diff = (x - y).abs().max()
+    cos_diff = 1 - 2 * (x * y).sum().item() / max((x * x + y * y).sum().item(), 1e-12)
+    assert diff < 5e-3
+    assert cos_diff < 5e-3
+
+
+def get_padding_offset(bsz, max_seq_len, seq_lens_this_time):
+    cum_offsets_now = paddle.cumsum(max_seq_len - seq_lens_this_time)
+    cum_offsets = paddle.zeros(shape=(bsz + 1), dtype="int32")
+    cum_offsets[1:] = cum_offsets_now
+    token_num = paddle.sum(seq_lens_this_time)
+    padding_offsets = paddle.zeros(shape=(token_num), dtype="int32")
+    cu_seqlens_q = paddle.zeros(shape=(bsz + 1), dtype="int32")
+    cu_seqlens_k = paddle.zeros(shape=(bsz + 1), dtype="int32")
+    for i in range(bsz):
+        seq_len_now = seq_lens_this_time[i]
+        cum_offset = cum_offsets[i]
+        for j in range(seq_len_now):
+            padding_offsets[i * max_seq_len - cum_offset + j] = cum_offset
+        cum_seq_len = (i + 1) * max_seq_len - cum_offsets[i + 1]
+        cu_seqlens_q[i + 1] = cum_seq_len
+        cu_seqlens_k[i + 1] = cum_seq_len
+    return padding_offsets, cum_offsets[:-1], cu_seqlens_q, cu_seqlens_k
+
+
+RUN_TIME = 1
+WARM_UP = 0
+BLOCK_SIZE = 64
+HEAD_DIM_QK = 576
+HEAD_DIM_V = 512
+PE_SIZE = 64
+MAX_LENGTH = 8192
+MAX_DEC_LEN = 1024
+NUM_Q_HEAD = 8
+NUM_KV_HEAD = 1
+dtype = "bfloat16"
+DRAFT_TOTAL_TOKEN_NUM = 1
+CAUSAL = True
+
+SPECULATE_DECODER = False
+
+
+def ref_attention(q_all, p_compressed_kv, compressed_kv, p_key_pe, key_pe, bsz, cache_length, softmax_scale):
+    q = q_all.reshape([bsz, DRAFT_TOTAL_TOKEN_NUM, NUM_Q_HEAD, HEAD_DIM_QK])
+    p_compressed_kv = p_compressed_kv.reshape((bsz, cache_length, HEAD_DIM_V))
+    compressed_kv = compressed_kv.reshape((bsz, DRAFT_TOTAL_TOKEN_NUM, HEAD_DIM_V))
+    p_key_pe = p_key_pe.reshape((bsz, cache_length, PE_SIZE))
+    key_pe = key_pe.reshape((bsz, DRAFT_TOTAL_TOKEN_NUM, PE_SIZE))
+
+    p_k = paddle.concat([p_compressed_kv, p_key_pe], axis=-1).reshape((bsz, cache_length, HEAD_DIM_QK))
+    k = paddle.concat([compressed_kv, key_pe], axis=-1).reshape((bsz, DRAFT_TOTAL_TOKEN_NUM, HEAD_DIM_QK))
+    k = paddle.concat([p_k, k], axis=1).transpose([0, 2, 1])
+    v = paddle.concat([p_compressed_kv, compressed_kv], axis=1).reshape(
+        (bsz, cache_length + DRAFT_TOTAL_TOKEN_NUM, HEAD_DIM_V)
+    )
+
+    out = paddle.zeros(shape=[bsz, DRAFT_TOTAL_TOKEN_NUM, NUM_Q_HEAD, HEAD_DIM_V], dtype=dtype)
+    for i in range(DRAFT_TOTAL_TOKEN_NUM):
+        for j in range(NUM_Q_HEAD):
+
+            query = q[:, i, j, :].reshape([-1, 1, HEAD_DIM_QK])
+            query = query.scale(softmax_scale)
+            cu_k = k[:, :, : cache_length + i + 1]
+            score = paddle.bmm(x=query, y=cu_k)
+            score = paddle.nn.functional.softmax(score, axis=-1)
+            cu_v = v[:, : cache_length + i + 1, :]
+            sub_out = paddle.bmm(score, cu_v).reshape([-1, 1, 1, HEAD_DIM_V])
+            out[:, i : i + 1, j : j + 1, :] = sub_out
+
+    out = out.reshape([-1, NUM_Q_HEAD, HEAD_DIM_V])
+    return out
+
+
+def fake_prefill(input_length, compressed_kv, key_pe, latent_cache_shape, block_tables):
+
+    seq_lens_enc = [
+        input_length,
+    ] * bsz
+    seq_lens_dec = [
+        0,
+    ] * bsz
+    seq_lens_this_time = [
+        input_length,
+    ] * bsz
+    max_enc_len_this_time = max(seq_lens_enc)
+    max_dec_len_this_time = max(seq_lens_dec)
+    max_enc_len_this_time = paddle.to_tensor([max_enc_len_this_time], "int32", place=paddle.CPUPlace())
+    max_dec_len_this_time = paddle.to_tensor([max_dec_len_this_time], "int32", place=paddle.CPUPlace())
+
+    seq_lens_encoder = paddle.to_tensor(seq_lens_enc, "int32")
+    seq_lens_this_time = paddle.to_tensor(seq_lens_this_time, "int32")
+    seq_lens_decoder = paddle.to_tensor(seq_lens_dec, "int32")
+    padding_offsets, cum_offsets, cu_seqlens_q, cu_seqlens_k = get_padding_offset(bsz, MAX_LENGTH, seq_lens_this_time)
+    latent_cache = paddle.zeros(shape=latent_cache_shape).astype(dtype)
+    # import pdb; pdb.set_trace()
+
+    paddlenlp_ops.prefill_mla_write_cache(
+        compressed_kv,
+        key_pe,
+        latent_cache,
+        seq_lens_encoder,
+        seq_lens_decoder,
+        padding_offsets,
+        cum_offsets,
+        block_tables,
+        "none",
+        MAX_LENGTH,
+    )
+
+    return latent_cache
+
+
+def test_append_c16_attention(cache_length, bsz):
+    # prepare input
+    seq_lens_enc = [
+        0,
+    ] * bsz
+    seq_lens_dec = [
+        cache_length,
+    ] * bsz
+    seq_lens_this_time = [
+        DRAFT_TOTAL_TOKEN_NUM,
+    ] * bsz
+    max_enc_len_this_time = max(seq_lens_enc)
+    max_dec_len_this_time = max(seq_lens_dec)
+    max_enc_len_this_time = paddle.to_tensor([max_enc_len_this_time], "int32", place=paddle.CPUPlace())
+    max_dec_len_this_time = paddle.to_tensor([max_dec_len_this_time], "int32", place=paddle.CPUPlace())
+    token_num = sum(seq_lens_this_time)
+    block_num_per_seq = div_up(MAX_LENGTH, BLOCK_SIZE)
+    max_block_num = block_num_per_seq * bsz
+    free_list = list(range(max_block_num - 1, -1, -1))
+    block_tables = paddle.ones(shape=(bsz, block_num_per_seq), dtype="int32") * (-1)
+    for i in range(bsz):
+        need_block_num = div_up(seq_lens_dec[i] + MAX_DEC_LEN, BLOCK_SIZE)
+        for j in range(need_block_num):
+            block_id = free_list.pop()
+            block_tables[i, j] = block_id
+    seq_lens_encoder = paddle.to_tensor(seq_lens_enc, "int32")
+    seq_lens_this_time = paddle.to_tensor(seq_lens_this_time, "int32")
+    seq_lens_decoder = paddle.to_tensor(seq_lens_dec, "int32")
+    padding_offsets, cum_offsets, cu_seqlens_q, cu_seqlens_k = get_padding_offset(bsz, MAX_LENGTH, seq_lens_this_time)
+    q_varlen_shape = [token_num, NUM_Q_HEAD * HEAD_DIM_QK]
+    latent_cache_shape = (
+        max_block_num,
+        NUM_KV_HEAD,
+        BLOCK_SIZE,
+        HEAD_DIM_QK,
+    )
+
+    (
+        encoder_batch_ids,
+        encoder_tile_ids_per_batch,
+        encoder_num_blocks,
+        kv_batch_ids,
+        kv_tile_ids_per_batch,
+        kv_num_blocks,
+        decoder_batch_ids,
+        decoder_tile_ids_per_batch,
+        decoder_num_blocks_device,
+        decoder_num_blocks,
+        decoder_chunk_size,
+        max_len_kv,
+    ) = paddlenlp_ops.get_block_shape_and_split_kv_block(
+        seq_lens_encoder,
+        seq_lens_decoder,
+        max_enc_len_this_time,
+        max_dec_len_this_time,
+        seq_lens_this_time,
+        cum_offsets,
+        NUM_Q_HEAD // NUM_KV_HEAD,
+        BLOCK_SIZE,
+        DRAFT_TOTAL_TOKEN_NUM,
+    )
+    softmax_scale = HEAD_DIM_QK ** (-0.5)
+
+    # fake_prefill to prepare latent_cache
+    p_compressed_kv_shape = [bsz * cache_length, NUM_KV_HEAD * HEAD_DIM_V]
+    p_key_pe_shape = [cache_length * bsz, NUM_KV_HEAD * PE_SIZE]
+    p_compressed_kv = paddle.randn(shape=p_compressed_kv_shape).astype(dtype)
+    p_key_pe = paddle.randn(shape=p_key_pe_shape).astype(dtype)
+    latent_cache = fake_prefill(cache_length, p_compressed_kv, p_key_pe, latent_cache_shape, block_tables)
+
+    # dec
+    query = paddle.randn(shape=q_varlen_shape).astype(dtype)
+    compressed_kv_shape = [token_num, NUM_KV_HEAD, HEAD_DIM_V]
+    key_pe_shape = [token_num, NUM_KV_HEAD, PE_SIZE]
+    compressed_kv = paddle.rand(shape=compressed_kv_shape).astype(dtype)
+    key_pe = paddle.rand(shape=key_pe_shape).astype(dtype)
+    paddlenlp_ops.decode_mla_write_cache(
+        compressed_kv,
+        key_pe,
+        latent_cache,
+        seq_lens_decoder,
+        seq_lens_encoder,
+        padding_offsets,
+        cum_offsets,
+        block_tables,
+        "none",
+        MAX_LENGTH,
+        SPECULATE_DECODER,
+    )
+
+    ref_out = ref_attention(query, p_compressed_kv, compressed_kv, p_key_pe, key_pe, bsz, cache_length, softmax_scale)
+    paddle.device.synchronize()
+    s_time = 0
+    for i in range(RUN_TIME + WARM_UP):
+        if i == WARM_UP:
+            s_time = time.time()
+        out = paddlenlp_ops.multi_head_latent_attention(
+            query,
+            latent_cache,
+            latent_cache,
+            seq_lens_encoder,
+            seq_lens_decoder,
+            seq_lens_this_time,
+            cu_seqlens_q,
+            padding_offsets,
+            cum_offsets,
+            block_tables,
+            encoder_batch_ids,
+            encoder_tile_ids_per_batch,
+            encoder_num_blocks,
+            kv_batch_ids,
+            kv_tile_ids_per_batch,
+            kv_num_blocks,
+            decoder_batch_ids,
+            decoder_tile_ids_per_batch,
+            decoder_num_blocks_device,
+            decoder_num_blocks,
+            decoder_chunk_size,
+            max_enc_len_this_time,
+            max_dec_len_this_time,
+            max_len_kv,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+            "bf16",
+            "none",  # cache_quant_type
+            HEAD_DIM_V,
+            MAX_LENGTH,
+            softmax_scale,
+            0.0,
+            0.0,
+            0.0,  # out_linear_in_scale
+            DRAFT_TOTAL_TOKEN_NUM,  # DRAFT_TOTAL_TOKEN_NUM
+            CAUSAL,
+            SPECULATE_DECODER,
+        )
+        paddle.device.synchronize()
+
+    e_time = time.time()
+    out = out.reshape([-1, NUM_Q_HEAD, HEAD_DIM_V])
+    alloc_diff(ref_out, out)
+    print(
+        "dec bsz:{}, num_q_head:{}, cache_length:{}, cost_time:{}ms".format(
+            bsz, NUM_Q_HEAD, cache_length, (e_time - s_time) / RUN_TIME * 1000
+        )
+    )
+
+
+if __name__ == "__main__":
+    print("CUDA_VERSION >=12.8 is required to compile and run mla benchmark")
+    for cache_length in [1023, 2047]:
+        for bsz in [1, 16, 32, 64, 128, 256]:
+            test_append_c16_attention(cache_length, bsz)

csrc/gpu/unittest/test_cutlass_block_gemm.py

@@ -0,0 +1,119 @@
+# 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 unittest
+
+import numpy as np
+import paddle
+from paddlenlp_ops import cutlass_fp8_fp8_half_block_gemm_fused as fp8_block_gemm_fused
+
+paddle.seed(1)
+
+
+def create_random_cuda_tensor(shape, dtype, mean: float = 0, std: float = 1):
+    return paddle.empty(shape, dtype=dtype).normal_(mean, std)
+
+
+class CutlassBlockGemmTest(unittest.TestCase):
+    def native_w8a8_block_fp8_matmul(self, A, B, As, Bs, block_size, output_dtype="float16"):
+        """This function performs matrix multiplication with block-wise quantization using native paddle.
+
+        It takes two input tensors `A` and `B` with scales `As` and `Bs`.
+        The output is returned in the specified `output_dtype`.
+        """
+
+        A = A.cast(paddle.float32)
+        B = B.cast(paddle.float32)
+        assert A.shape[-1] == B.shape[-1]
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+        assert (A.shape[-1] + block_k - 1) // block_k == As.shape[-1]
+        assert A.shape[:-1] == As.shape[:-1]
+
+        M = A.numel() // A.shape[-1]
+        N, K = B.shape
+        origin_C_shape = A.shape[:-1] + [N]
+        A = A.reshape([M, A.shape[-1]])
+        As = As.reshape([M, As.shape[-1]])
+        n_tiles = (N + block_n - 1) // block_n
+        k_tiles = (K + block_k - 1) // block_k
+        assert n_tiles == Bs.shape[0]
+        assert k_tiles == Bs.shape[1]
+
+        C_shape = [M, N]
+        C = paddle.zeros(C_shape, dtype=paddle.float32)
+        A_tiles = [A[:, i * block_k : min((i + 1) * block_k, K)] for i in range(k_tiles)]
+        B_tiles = [
+            [
+                B[
+                    j * block_n : min((j + 1) * block_n, N),
+                    i * block_k : min((i + 1) * block_k, K),
+                ]
+                for i in range(k_tiles)
+            ]
+            for j in range(n_tiles)
+        ]
+        C_tiles = [C[:, j * block_n : min((j + 1) * block_n, N)] for j in range(n_tiles)]
+        As_tiles = [As[:, i : i + 1] for i in range(k_tiles)]
+
+        for i in range(k_tiles):
+            for j in range(n_tiles):
+                a = A_tiles[i]
+                b = B_tiles[j][i]
+                c = C_tiles[j]
+                s = As_tiles[i] * Bs[j][i]
+                c[:, :] += paddle.matmul(a, b.t()) * s
+        C = C.reshape(origin_C_shape).cast(output_dtype)
+        return C
+
+    def test_cutlass_fp8_block_gemm_fused(self):
+        M = 32
+        K = 7168
+        N = 1536
+        block_size = [128, 128]
+        output_dtype = "float16"
+
+        scale_k = (K + 128 - 1) // 128
+        scale_n = (N + 128 - 1) // 128
+        A = (
+            create_random_cuda_tensor([M, K], "float32", mean=0, std=0.1)
+            .clip(min=-448.0, max=448.0)
+            .cast(paddle.float8_e4m3fn)
+        )
+        B = (
+            create_random_cuda_tensor([N, K], "float32", mean=0, std=0.1)
+            .clip(min=-448.0, max=448.0)
+            .cast(paddle.float8_e4m3fn)
+        )
+        As = paddle.rand([M, scale_k], dtype=paddle.float32).clip(min=0.001)
+        Bs = paddle.rand([scale_n, scale_k], dtype=paddle.float32).clip(min=0.001)
+
+        ref_out = self.native_w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype)
+        out = fp8_block_gemm_fused(
+            A,
+            B,
+            As.t(),
+            Bs,
+            bias=None,
+            transpose_x=False,
+            transpose_y=True,
+            output_dtype=output_dtype,
+            act="identity",
+        )
+
+        np.testing.assert_allclose(ref_out.numpy(), out.numpy(), rtol=1e-3, atol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()

csrc/gpu/unittest/test_fused_moe.py

@@ -0,0 +1,210 @@
+# 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 unittest
+
+import numpy as np
+from paddlenlp_ops import (
+    fused_expert_moe,
+    moe_expert_dispatch,
+    moe_expert_ffn,
+    moe_expert_reduce,
+)
+
+import paddle
+import paddle.nn.functional as F
+from paddle import nn
+from paddle.incubate.nn.functional import swiglu
+
+paddle.seed(42)
+np.random.seed(42)
+
+class Expert(nn.Layer):
+    def __init__(self, d_model, d_feedforward):
+        super().__init__()
+        self.fc1 = nn.Linear(d_model, d_feedforward * 2)
+        self.fc2 = nn.Linear(d_feedforward, d_model)
+
+    def forward(self, x, idx):
+        x = self.fc1(x)
+        x = swiglu(x)
+        return self.fc2(x)
+
+class TestSimpleFusedMoe(unittest.TestCase):
+    def setUp(self):
+        self.set_config()
+        self.init_parameters()
+        self.prepare_data()
+
+
+    def set_config(self):
+        self.dtype = "bfloat16"
+        self.batch_size = 16
+        self.seq_len = 128
+        self.num_experts = 4
+        self.d_model = 64
+        self.d_feedforward = 128
+        self.top_k = 2
+        self.rtol = 1e-2
+        self.atol = 1e-2
+
+        paddle.set_default_dtype(self.dtype)
+
+    def init_parameters(self):
+        # 创建专家层
+        self.experts = nn.LayerList([
+            Expert(self.d_model, self.d_feedforward)
+            for _ in range(self.num_experts)
+        ])
+
+        # 初始化门控权重
+        self.gate = nn.Linear(self.d_model, self.num_experts)
+        self.gate_weight = self.gate.weight.cast("float32")
+
+    def prepare_data(self):
+        """准备输入数据"""
+        self.x = paddle.randn(
+            [self.batch_size, self.seq_len, self.d_model],
+            dtype=self.dtype
+        )
+
+        self.w0 = paddle.stack(
+            [e.fc1.weight for e in self.experts], axis=0
+        ).astype(self.dtype)
+        self.b0 = paddle.stack(
+            [e.fc1.bias for e in self.experts], axis=0
+        ).reshape([self.num_experts, 1, -1]).astype(self.dtype)
+
+        self.w1 = paddle.stack(
+            [e.fc2.weight for e in self.experts], axis=0
+        ).astype(self.dtype)
+        self.b1 = paddle.stack(
+            [e.fc2.bias for e in self.experts], axis=0
+        ).reshape([self.num_experts, 1, -1]).astype(self.dtype)
+
+    def baseline_forward(self, hidden_states):
+        """（逐个专家计算）"""
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.reshape([-1, hidden_dim])
+
+        # 路由计算
+        logits = paddle.matmul(hidden_states.cast("float32"), self.gate_weight)
+        weights = F.softmax(logits, axis=-1)
+        routing_weights, selected_experts = paddle.topk(weights, self.top_k, axis=-1)
+        # 结果累加
+        final_hidden_states = paddle.zeros_like(hidden_states)
+
+        expert_mask = paddle.transpose(
+            F.one_hot(selected_experts, num_classes=self.num_experts), [2, 1, 0]
+        )
+
+        for expert_id in range(self.num_experts):
+            expert_layer = self.experts[expert_id]
+            idx, top_x = paddle.where(expert_mask[expert_id])
+
+            current_state = paddle.index_select(
+                hidden_states, top_x, axis=0
+            ).reshape([-1, hidden_dim])
+            current_hidden_states = (
+                expert_layer(current_state, expert_id)
+                * routing_weights[top_x, idx]
+            )
+            paddle.index_add_(
+                x=final_hidden_states,
+                index=top_x.squeeze(),
+                axis=0,
+                value=current_hidden_states.to(hidden_states.dtype),
+            )
+        final_hidden_states = paddle.reshape(
+            final_hidden_states, [batch_size, seq_len, hidden_dim]
+        )
+        return final_hidden_states
+
+    def fused_forward(self, x):
+        """测试融合实现"""
+        return fused_expert_moe(
+            x,
+            self.gate_weight,
+            self.w0,
+            self.w1,
+            self.b0,
+            None,
+            self.b1,
+            None,
+            "None",
+            self.top_k,
+            False,
+            False
+        )
+
+    def split_forward(self, hidden_states):
+        '''测试拆分实现'''
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.reshape([-1, hidden_dim])
+
+        # 路由计算
+        logits = paddle.matmul(hidden_states.cast("float32"), self.gate_weight)
+        scores = F.softmax(logits, axis=-1)
+        (
+            permute_input,
+            token_nums_per_expert,
+            permute_indices_per_token,
+            top_k_weights,
+            top_k_indices,
+        ) = moe_expert_dispatch(hidden_states, scores, self.top_k, False, topk_only_mode=True)
+
+        ffn_out = moe_expert_ffn(
+            permute_input,
+            token_nums_per_expert,
+            self.w0,
+            self.w1,
+            self.b0,
+            None,
+            None,
+            "none",
+        )
+        output = moe_expert_reduce(
+                ffn_out,
+                top_k_weights,
+                permute_indices_per_token,
+                top_k_indices,
+                None,
+                norm_topk_prob=False,
+                routed_scaling_factor=1.0,
+            )
+        output = paddle.reshape(output, [batch_size, seq_len, hidden_dim])
+        return output
+
+    def test_consistency(self):
+        """测试一致性"""
+        base_out = self.baseline_forward(self.x)
+        fused_out = self.fused_forward(self.x)
+        split_out = self.split_forward(self.x)
+
+        np.testing.assert_allclose(
+            base_out.cast("float32").numpy().astype("float32"),
+            fused_out.cast("float32").numpy().astype("float32"),
+            rtol=self.rtol,
+            atol=self.atol
+        )
+        np.testing.assert_allclose(
+            base_out.cast("float32").numpy().astype("float32"),
+            split_out.cast("float32").numpy().astype("float32"),
+            rtol=self.rtol,
+            atol=self.atol
+        )
+        print("Test passed! Outputs match within tolerance.")
+
+if __name__ == "__main__":
+    unittest.main()

csrc/setup.py

@@ -1,73 +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.
-
-""" setup for EfficentLLM """
-
-import os
-
-from setuptools import find_packages, setup
-
-description = "Paddlenlp_ops : inference framework implemented based on PaddlePaddle"
-VERSION = "0.0.0"
-
-
-def read(file: str):
-    """
-    read file and return content
-    """
-    current_dir = os.path.dirname(__file__)
-    path = os.path.join(current_dir, file)
-    with open(path, "r", encoding="utf-8") as f:
-        content = f.read().strip()
-    return content
-
-
-def read_version():
-    """
-    read version and return content
-    """
-    return VERSION
-
-
-def read_readme():
-    """
-    read README.md and return content
-    """
-    return read("README.md")
-
-
-setup(
-    name="paddlenlp_ops",
-    packages=find_packages(),
-    version="0.0.0",
-    author="Paddle Infernce Team",
-    author_email="paddle-inference@baidu.com",
-    description=description,
-    long_description=read_readme(),
-    long_description_content_type="text/markdown",
-    url="",
-    python_requires=">=3.8",
-    package_dir={"paddlenlp_ops": "paddlenlp_ops/"},
-    package_data={"paddlenlp_ops": ["sm70/*", "sm75/*", "sm80/*", "sm86/*", "sm89/*", "sm90/*"]},
-    include_package_data=True,
-    classifiers=[
-        "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.9",
-        "Programming Language :: Python :: 3.10",
-        "License :: OSI Approved :: Apache Software License",
-        "Operating System :: OS Independent",
-    ],
-    license="Apache 2.0",
-)

csrc/setup_cuda.py

@@ -80,7 +80,7 @@ def get_gencode_flags():
 
 
 gencode_flags = get_gencode_flags()
-library_path = os.environ.get("LD_LIBRARY_PATH", "/usr/local/cuda/lib64")
+library_path = [os.environ.get("LD_LIBRARY_PATH", "/usr/local/cuda/lib64")]
 
 sources = [
     "./gpu/save_with_output.cc",
@@ -105,39 +105,48 @@ sources = [
     "./gpu/step.cu",
     "./gpu/quant_int8.cu",
     "./gpu/dequant_int8.cu",
-    "./gpu/get_position_ids.cu",
+    "./gpu/group_quant.cu",
+    "./gpu/moe/preprocess_for_moe.cu",
+    "./gpu/get_position_ids_and_mask_encoder_batch.cu",
     "./gpu/fused_rotary_position_encoding.cu",
     "./gpu/flash_attn_bwd.cc",
     "./gpu/tune_cublaslt_gemm.cu",
     "./gpu/sample_kernels/top_p_sampling_reject.cu",
     "./gpu/update_inputs_v2.cu",
+    "./gpu/noaux_tc.cu",
     "./gpu/set_preids_token_penalty_multi_scores.cu",
     "./gpu/speculate_decoding_kernels/ngram_match.cc",
     "./gpu/speculate_decoding_kernels/speculate_save_output.cc",
     "./gpu/speculate_decoding_kernels/speculate_get_output.cc",
 ]
 sources += find_end_files("./gpu/speculate_decoding_kernels", ".cu")
+sources += find_end_files("./gpu/moe/fused_moe/cutlass_kernels/moe_gemm/", ".cu")
+sources += find_end_files("./gpu/moe/fused_moe/", ".cu")
 
 nvcc_compile_args = gencode_flags
 update_git_submodule()
 nvcc_compile_args += [
     "-O3",
+    "-DNDEBUG",
     "-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__",
-    "-Igpu",
-    "-Igpu/cutlass_kernels",
-    "-Igpu/fp8_gemm_with_cutlass",
-    "-Igpu/cutlass_kernels/fp8_gemm_fused/autogen",
-    "-Ithird_party/cutlass/include",
-    "-Ithird_party/cutlass/tools/util/include",
-    "-Ithird_party/nlohmann_json/single_include",
-    "-Igpu/sample_kernels",
 ]
 
+include_dirs = [
+    "./gpu",
+    "./gpu/cutlass_kernels",
+    "./gpu/fp8_gemm_with_cutlass",
+    "./gpu/cutlass_kernels/fp8_gemm_fused/autogen",
+    "./third_party/cutlass/include",
+    "./third_party/cutlass/tools/util/include",
+    "./third_party/nlohmann_json/single_include",
+    "./gpu/sample_kernels",
+    "./gpu/moe/fused_moe",
+]
 cc = get_sm_version()
 cuda_version = float(paddle.version.cuda())
 
@@ -154,6 +163,7 @@ fp8_auto_gen_directory = "gpu/cutlass_kernels/fp8_gemm_fused/autogen"
 if os.path.isdir(fp8_auto_gen_directory):
     shutil.rmtree(fp8_auto_gen_directory)
 
+
 if cc == 89 and cuda_version >= 12.4:
     os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels.py --cuda_arch 89")
     os.system("python utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels.py --cuda_arch 89")
@@ -164,24 +174,55 @@ if cc == 89 and cuda_version >= 12.4:
         "gpu/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu",
     ]
 
+if cc >= 80 and cuda_version >= 12.4:
+    nvcc_compile_args += [
+        "-std=c++17",
+        "--use_fast_math",
+        "--threads=8",
+        "-D_GLIBCXX_USE_CXX11_ABI=1",
+    ]
+    sources += ["./gpu/sage_attn_kernels/sageattn_fused.cu"]
+    if cc >= 80 and cc < 89:
+        sources += [
+            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f16_kernel_sm80.cu"
+        ]
+        nvcc_compile_args += ["-gencode", f"arch=compute_80,code=compute_80"]
+    elif cc >= 89 and cc < 90:
+        sources += [
+            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm89.cu"
+        ]
+        nvcc_compile_args += ["-gencode", f"arch=compute_89,code=compute_89"]
+    elif cc >= 90:
+        sources += [
+            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_kernel_sm90.cu",
+            "./gpu/sage_attn_kernels/sageattn_qk_int_sv_f8_dsk_kernel_sm90.cu"
+        ]
+        nvcc_compile_args += ["-gencode", f"arch=compute_90a,code=compute_90a"]
+
 if cc >= 90 and cuda_version >= 12.0:
-    nvcc_compile_args += ["-DNDEBUG"]
     os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py --cuda_arch 90")
+    os.system("python utils/auto_gen_fp8_fp8_gemm_fused_kernels_ptr_scale_sm90.py --cuda_arch 90")
     os.system("python utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels_sm90.py --cuda_arch 90")
+    os.system("python utils/auto_gen_fp8_fp8_block_gemm_fused_kernels_sm90.py --cuda_arch 90")
     sources += find_end_files(fp8_auto_gen_directory, ".cu")
     sources += [
         "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm.cu",
         "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_cuda_core_gemm.cu",
         "gpu/fp8_gemm_with_cutlass/fp8_fp8_fp8_dual_gemm.cu",
+        "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_block_gemm.cu",
+        "gpu/fp8_gemm_with_cutlass/fp8_fp8_half_gemm_ptr_scale.cu",
     ]
+    sources += find_end_files("./gpu/mla_attn", ".cu")
 
 ops_name = f"paddlenlp_ops_{sm_version}" if sm_version != 0 else "paddlenlp_ops"
+
 setup(
     name=ops_name,
     ext_modules=CUDAExtension(
         sources=sources,
-        extra_compile_args={"cxx": ["-O3"], "nvcc": nvcc_compile_args},
+        extra_compile_args={"cxx": ["-O3", "-fopenmp", "-lgomp", "-std=c++17", "-DENABLE_BF16"], "nvcc": nvcc_compile_args},
         libraries=["cublasLt"],
-        library_dirs=[library_path],
+        library_dirs=library_path,
+        include_dirs=include_dirs,
     ),
 )

csrc/setup_hip.py

@@ -12,15 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from paddle.utils.cpp_extension import CUDAExtension, setup
 import subprocess
 
+from paddle.utils.cpp_extension import CUDAExtension, setup
+
+
 def update_git_submodule():
     try:
         subprocess.run(["git", "submodule", "update", "--init"], check=True)
     except subprocess.CalledProcessError as e:
         print(f"Error occurred while updating git submodule: {str(e)}")
         raise
+
+
 update_git_submodule()
 setup(
     name="paddlenlp_ops",
@@ -42,7 +46,6 @@ setup(
             "./gpu/rebuild_padding_v2.cu",
             "./gpu/set_value_by_flags_v2.cu",
             "./gpu/stop_generation_multi_ends_v2.cu",
-            "./gpu/update_inputs.cu",
             "./gpu/get_output.cc",
             "./gpu/save_with_output_msg.cc",
             "./gpu/write_int8_cache_kv.cu",

csrc/tools/build_wheel.sh

@@ -92,11 +92,11 @@ function create_directories(){
 """ setup for PaddlenlpOps """
 
 import os
+from datetime import datetime
 
 from setuptools import find_packages, setup
 
 description = "Paddlenlp_ops : inference framework implemented based on PaddlePaddle"
-VERSION = "0.0.0"
 
 
 def read(file: str):
@@ -114,12 +114,16 @@ def read_version():
     """
     read version and return content
     """
-    return VERSION
+    __version__ = "3.0.0b4.post"
+    formatted_date = datetime.now().date().strftime("%Y%m%d")
+    __version__ = __version__.replace(".post", ".post{}".format(formatted_date))
+    
+    return __version__
 
 setup(
     name="paddlenlp_ops",
     packages=find_packages(),
-    version="0.0.0",
+    version=read_version(),
     author="Paddle Infernce Team",
     author_email="paddle-inference@baidu.com",
     description=description,

csrc/tools/tune_fp8_gemm.sh

@@ -19,17 +19,33 @@ export FLAGS_use_cutlass_device_best_config_path=tune
 
 # 根据不同模型，调整gemm参数，其中n和k的的值需要成对出现，具体逻辑可参考相关.py文件。
 # llama3-8B 单GEMM调优参数设置
-nohup python ./utils/tune_cutlass_fp8_gemm.py \
-        --m_min 32 \
-        --m_max 32768 \
-        --n 6144 4096 4096 \
-        --k 4096 4096 14336 \
-        >  tune_fp8_gemm.log 2>&1 &
+# python ../utils/tune_cutlass_fp8_gemm.py \
+#         --m_min 32 \
+#         --m_max 32768 \
+#         --n 6144 4096 4096 \
+#         --k 4096 4096 14336 \
+#         >  tune_fp8_gemm.log 2>&1
 
 # llama3-8B 双GEMM调优参数设置
-# nohup python ./utils/tune_cutlass_fp8_dual_gemm.py \
+# python ../utils/tune_cutlass_fp8_dual_gemm.py \
 #         --m_min 32 \
 #         --m_max 32768 \
 #         --n 14336 \
 #         --k 4096 \
-#         >  tune_fp8_dual_gemm.log 2>&1 &
+#         >  tune_fp8_dual_gemm.log 2>&1
+
+# deepseek-v3 tp16调优
+# python ../utils/tune_cutlass_fp8_block_gemm.py \
+#         --m_min 32 \
+#         --m_max 32768 \
+#         --n 1536 1536 576 2048 4096 512 7168 7168 2304 7168 256 7168 \
+#         --k 7168 1536 7168 512 1536 1536 4096 1024 7168 1152 7168 128 \
+#         >  deepseekv3_fp8_config.log 2>&1 
+
+# fp8 gemm with scale tensor on sm90
+# python ../utils/tune_cutlass_fp8_gemm_ptr_scale.py \
+#         --m_min 32 \
+#         --m_max 32768 \
+#         --n 6144 4096 4096 \
+#         --k 4096 4096 14336 \
+#         >  tune_fp8_gemm_ptr_scale.log 2>&1

csrc/utils/auto_gen_fp8_fp8_block_gemm_fused_kernels_sm90.py

@@ -0,0 +1,582 @@
+# 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 argparse
+import copy
+import os
+import re
+
+
+def get_candidate_tiles():
+    cta_shape = [
+        ("<_128, _128, _128>"),
+        # ("<_256, _128, _128>"),
+    ]
+    cluster_shape = [
+        ("<_1, _1, _1>"),
+        ("<_2, _1, _1>"),
+        ("<_1, _2, _1>"),
+        ("<_2, _2, _1>"),
+        # ("<_1, _8, _1>"),
+        # ("<_8, _1, _1>"),
+    ]
+    base_configs = [(x, y) for x in cta_shape for y in cluster_shape]
+
+    return base_configs
+
+
+def get_candidate_configs(sm):
+    tiles = get_candidate_tiles()
+    candidate_configs = list()
+
+    hasbias = ("false", "true")
+    KernelSchedule = ("KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<1>",)
+    EpilogueSchedule = ("TmaWarpSpecializedCooperative",)
+    TileSchedule = ("PersistentScheduler", "StreamKScheduler")
+    for act_tag in [
+        ("noact", "Identity"),
+        # ("relu", "ReLu"),
+        # ("gelu", "GELU"),
+    ]:
+        candidate_configs.extend([(hasbias, act_tag, tiles, KernelSchedule, EpilogueSchedule, TileSchedule)])
+    return candidate_configs
+
+
+def get_shape_str(tile_shape):
+    blocks, clusters = [s.replace(" ", "").strip("<>").split(",") for s in tile_shape]
+    blocks = [elem.strip("_") for elem in blocks]
+    clusters = [elem.strip("_") for elem in clusters]
+    return blocks, clusters
+
+
+def check_config_valid(tile_shape, kernel_schedule, epilogue_schedule, tile_schedule):
+    blocks, clusters = get_shape_str(tile_shape)
+    if int(blocks[0]) < 128 and "Cooperative" in kernel_schedule:
+        return False
+    # if "Cooperative" in kernel_schedule and "Cooperative" not in epilogue_schedule:
+    #     return False
+    # if (
+    #     tile_shape[0] == "<_256, _128, _128>"
+    #     and "Cooperative" not in kernel_schedule
+    #     and "Cooperative" not in epilogue_schedule
+    # ):
+    #     return False
+    # flag1 = (int(blocks[0]) == 64 and kernel_schedule == "KernelTmaWarpSpecializedPingpongFP8FastAccum" and epilogue_schedule == "TmaWarpSpecialized")
+    # flag2 = (int(blocks[0]) != 64 and kernel_schedule == "KernelTmaWarpSpecializedCooperativeFP8FastAccum" and epilogue_schedule == "TmaWarpSpecializedCooperative")
+    # if not (flag1 or flag2):
+    #     return False
+    return True
+
+
+# this is a file's header part
+CommonHead = """// Generated by auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py - Do not edit.
+
+#pragma once
+
+#include "fp8_gemm_fused/fuse_block_gemm_act_template_3x.h"
+
+"""
+
+GemmDeclare = """
+template<>
+bool dispatch_fuse_block_gemm_c3x<phi::dtype::{input_type}, phi::dtype::{output_type},
+                                 {hasbias},
+                                 cutlass::epilogue::thread::{Activation},
+                                 Shape{TileShape},
+                                 Shape{ClusterShape},
+                                 cutlass::gemm::{KernelSchedule},
+                                 cutlass::epilogue::{EpilogueSchedule},
+                                 cutlass::gemm::{TileSchedule},
+                                 {SM}
+                                 >(GemmEpilogueAllParams);
+
+"""
+
+
+LaunchGemmHead = """
+#pragma once
+
+#include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
+
+"""
+
+LaunchGemmDeclare = """
+bool launch_block_gemm_kernel_sm{sm}_{gemm_config}(const int type_id, GemmEpilogueAllParams params);
+"""
+
+LaunchGemmPart0 = """
+#pragma once
+
+#include "launch_block_gemm_kernel_sm{sm}.h"
+
+bool launch_block_gemm_kernel_sm{sm}_{gemm_config}(const int type_id, GemmEpilogueAllParams params){
+    switch (type_id) {
+"""
+
+LaunchGemmPart1 = """
+        case {type_id}:
+            return dispatch_fuse_block_gemm_c3x<phi::dtype::{input_type}, phi::dtype::{output_type},
+                                 {hasbias},
+                                 cutlass::epilogue::thread::{Activation},
+                                 Shape{TileShape},
+                                 Shape{ClusterShape},
+                                 cutlass::gemm::{KernelSchedule},
+                                 cutlass::epilogue::{EpilogueSchedule},
+                                 cutlass::gemm::{TileSchedule},
+                                 {SM}
+                                 >(params);
+            break;
+"""
+
+LaunchGemmPart2 = """
+        default:
+            throw std::runtime_error("cutlass gemm config is invalid.");
+            break;
+    }
+    return false;
+}
+"""
+
+
+code_part0 = """// Generated by auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py - Do not edit.
+
+#include <map>
+#include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
+#include "launch_block_gemm_kernel_sm{sm}.h"
+
+COMMON_DECLARE_string(use_cutlass_device_best_config_path);
+
+std::map<std::string, int> block_gemm_type_map{"""
+
+code_part1 = """
+    {"{input_type}_{output_type}_{hasbias}_{act_tag}",   {type_id}}, """
+
+code_part2 = """
+};
+
+std::map<std::string, int> block_gemm_config_map{
+"""
+
+code_part3 = """    {"{TileShape}, {ClusterShape}, {kernel_schedule}, {epilogue_schedule}, {tile_schedule}", {tile_id}},
+"""
+
+code_part4 = """};
+
+bool launch_block_gemm_kernel(const int type_id, const int kernel_id, GemmEpilogueAllParams params){
+    switch (kernel_id) {"""
+
+code_part5 = """
+        case {tile_id}:
+            return launch_block_gemm_kernel_sm{sm}_{gemm_config}(type_id, params);
+            break;
+"""
+
+code_part6 = """
+        default:
+            throw std::runtime_error("fp8 gemm_fused Config is invalid.");
+            break;
+    }
+    return false;
+}
+
+template <typename T>
+T get_relative_best(nlohmann::json* json_data,
+                    const std::string& target_key,
+                    const int& m,
+                    const int& n,
+                    const int& k) {
+    if (json_data->contains(target_key)) {
+        return json_data->at(target_key);
+    } else {
+        if (k > 3 * n){
+            return "<_128, _128, _128>, <_1, _2, _1>, KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<1>, TmaWarpSpecializedCooperative, StreamKScheduler";
+        }else{
+            return "<_128, _128, _128>, <_1, _2, _1>, KernelTmaWarpSpecializedCooperativeFP8BlockScaledSubGroupMAccum<1>, TmaWarpSpecializedCooperative, PersistentScheduler";
+        }
+
+    }
+}
+
+bool fp8_fp8_block_gemm_scale_bias_act(GemmEpilogueAllParams params) {
+  if (block_gemm_type_map.find(params.fuse_gemm_config) == block_gemm_type_map.end()) {
+    throw std::runtime_error("fp8 gemm_fused config is invalid.");
+  }
+
+  int type_id = block_gemm_type_map[params.fuse_gemm_config];
+  int M = (params.M + 31) / 32 * 32;
+  int N = params.N;
+  int K = params.K;
+
+  int kernel_id;
+  std::string mnk_string = "block_gemm_sm90<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
+  std::string best_config;
+  CutlassGemmConfigMannager& best_config_mannager = CutlassGemmConfigMannager::getInstance();
+  char *config_file_path_c_str = getenv("FLAGS_use_cutlass_device_best_config_path");
+  std::string config_file_path = config_file_path_c_str == nullptr ? "" : std::string(config_file_path_c_str);
+  if(config_file_path == "tune"){ // tune kernel
+    int warm_up_times = 5;
+    int tune_times = 10;
+    std::string best_kernel_id = "";
+    float duratation = 1000000.f;
+    // tune all kernel_id kernels
+    for(const auto& config_pair : block_gemm_config_map){
+        std::cout << "Running tune kernel: " << config_pair.first<< std::endl;
+        bool is_valid = true;
+        // warm up
+        for(int num_time = 0; num_time < warm_up_times; ++num_time){
+            if(!launch_block_gemm_kernel(type_id, config_pair.second, params)){
+                is_valid = false;
+                break;
+            }
+        }
+        if(!is_valid){
+            continue;
+        }
+        cudaEvent_t start, stop;
+        cudaEventCreate(&start);
+        cudaEventCreate(&stop);
+        cudaStreamSynchronize(params.stream);
+        cudaEventRecord(start, params.stream);
+        for(int num_time = 0; num_time < tune_times; ++num_time){
+            if(!launch_block_gemm_kernel(type_id, config_pair.second, params)){
+                is_valid = false;
+                break;
+            };
+        }
+        cudaEventRecord(stop, params.stream);
+        cudaEventSynchronize(stop);
+        float elapsedTime;
+        if(is_valid){
+            cudaEventElapsedTime(&elapsedTime, start, stop);
+        } else {
+            continue;
+        }
+        cudaEventDestroy(start);
+        cudaEventDestroy(stop);
+        if(elapsedTime < duratation){
+            best_kernel_id = config_pair.first;
+            duratation = elapsedTime;
+        }
+    }
+
+    nlohmann::json new_json;
+    new_json[mnk_string] = best_kernel_id;
+    best_config_mannager.up_date_configs(new_json);
+    std::cout <<"Gemm tune result for " << mnk_string<< ": best config is: "<< best_kernel_id << std::endl;
+    return true;
+  } else { // run kernel
+    nlohmann::json* config_json = new nlohmann::json();
+    if (config_file_path != "" && config_file_path != "default") {
+        config_json = best_config_mannager.get_gemm_best_configs(config_file_path);
+    }
+
+    best_config = get_relative_best<std::string>(config_json, mnk_string, M, N, K);
+
+    if (block_gemm_config_map.find(best_config) == block_gemm_config_map.end()) {
+        throw std::runtime_error("This config'kernel not be generate, please check auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py and re-generate.");
+    } else {
+        kernel_id = block_gemm_config_map[best_config];
+    }
+    return launch_block_gemm_kernel(type_id, kernel_id, params);
+  }
+}
+
+"""
+
+
+def SubstituteTemplate(template, values_base):
+    values = copy.deepcopy(values_base)
+    if values.get("KernelSchedule") is not None and "Auto" in values["KernelSchedule"]:
+        values["KernelSchedule"] = "collective::" + values["KernelSchedule"]
+    if values.get("EpilogueSchedule") is not None and "Auto" in values["EpilogueSchedule"]:
+        values["EpilogueSchedule"] = "collective::" + values["EpilogueSchedule"]
+    text = template
+    changed = True
+    while changed:
+        changed = False
+        for key, value in values.items():
+            regex = f"\\{{{key}\\}}"
+            newtext = re.sub(regex, value, text)
+            if newtext != text:
+                changed = True
+            text = newtext
+    return text
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="The argument for generating the generic_mixed_gemm_kernelLauncher instance."
+    )
+    parser.add_argument(
+        "--cuda_arch",
+        type=str,
+        nargs="+",
+        default=["90"],
+        help="The CUDA architecture to be generated.",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+# generate source .cu
+def generate_source_cu(
+    inputs_type: (str),
+    outputs_type: (str),
+    hasbiases: (str),
+    act_tag: (str),
+    tiles: (str),
+    KernelSchedule: (str),
+    EpilogueSchedule: (str),
+    TileSchedule: (str),
+    sm: str,
+):
+    all_code = CommonHead
+
+    for input_type in inputs_type:
+        for output_type in outputs_type:
+            for hasbias in hasbiases:
+                for tile_config in tiles:
+                    for kernel_schedule in KernelSchedule:
+                        for epilogue_schedule in EpilogueSchedule:
+                            for tile_schedule in TileSchedule:
+                                if not check_config_valid(
+                                    tile_config, kernel_schedule, epilogue_schedule, tile_schedule
+                                ):
+                                    continue
+                                value_dict = {
+                                    "input_type": input_type,
+                                    "output_type": output_type,
+                                    "hasbias": hasbias,
+                                    "Activation": act_tag[1],
+                                    "TileShape": tile_config[0],
+                                    "ClusterShape": tile_config[1],
+                                    "KernelSchedule": kernel_schedule,
+                                    "EpilogueSchedule": epilogue_schedule,
+                                    "TileSchedule": tile_schedule,
+                                    "SM": sm,
+                                    "sm": sm[-2:],
+                                }
+                                all_code += SubstituteTemplate(GemmDeclare, value_dict)
+
+    return all_code
+
+
+# generate gemm launch .cu
+def generate_launch_gemm_cus(
+    generate_dir: (str), inputs_type: (str), outputs_type: (str), fuse_gemm_configs: tuple, sm: str
+):
+    act_tags = [single_config[1] for single_config in fuse_gemm_configs]
+
+    single_config = fuse_gemm_configs[0]
+    hasbiases: (str) = single_config[0]
+    tiles: (str) = single_config[2]
+    KernelSchedule: (str) = single_config[3]
+    EpilogueSchedule: (str) = single_config[4]
+    TileSchedule: (str) = single_config[5]
+    code_map = {}
+    head_path = os.path.join(generate_dir, f"launch_block_gemm_kernel_sm{sm[-2:]}.h")
+    head_all_code = LaunchGemmHead
+    for tile_config in tiles:
+        blocks, clusters = get_shape_str(tile_config)
+        gemm_config_str_0 = f"tile{blocks[0]}x{blocks[1]}x{blocks[2]}_cluster{clusters[0]}x{clusters[1]}x{clusters[2]}"
+        for kernel_schedule in KernelSchedule:
+            gemm_config_str_1 = gemm_config_str_0 + f"_{kernel_schedule}"
+            for epilogue_schedule in EpilogueSchedule:
+                gemm_config_str_2 = gemm_config_str_1 + f"_{epilogue_schedule}"
+                for tile_schedule in TileSchedule:
+                    if not check_config_valid(tile_config, kernel_schedule, epilogue_schedule, tile_schedule):
+                        continue
+                    gemm_config_str = gemm_config_str_2 + f"_{tile_schedule}"
+                    value_dict = {
+                        "sm": sm[-2:],
+                        "gemm_config": gemm_config_str.replace("<", "").replace(">", ""),
+                    }
+                    head_all_code += SubstituteTemplate(LaunchGemmDeclare, value_dict)
+    os.makedirs(generate_dir, exist_ok=True)
+    with open(head_path, "w") as f:
+        f.write(head_all_code)
+        f.close()
+
+    for tile_shape in tiles:
+        blocks, clusters = get_shape_str(tile_shape)
+        gemm_config_str_0 = f"tile{blocks[0]}x{blocks[1]}x{blocks[2]}_cluster{clusters[0]}x{clusters[1]}x{clusters[2]}"
+        for kernel_schedule in KernelSchedule:
+            gemm_config_str_1 = gemm_config_str_0 + f"_{kernel_schedule}"
+            for epilogue_schedule in EpilogueSchedule:
+                gemm_config_str_2 = gemm_config_str_1 + f"_{epilogue_schedule}"
+                for tile_schedule in TileSchedule:
+                    if not check_config_valid(tile_shape, kernel_schedule, epilogue_schedule, tile_schedule):
+                        continue
+                    gemm_config_str = gemm_config_str_2 + f"_{tile_schedule}"
+                    value_dict = {
+                        "sm": sm[-2:],
+                        "gemm_config": gemm_config_str.replace("<", "").replace(">", ""),
+                    }
+                    source_all_code = SubstituteTemplate(LaunchGemmPart0, value_dict)
+                    type_id = 0
+                    for input_type in inputs_type:
+                        for output_type in outputs_type:
+                            for act_tag in act_tags:
+                                for hasbias in hasbiases:
+                                    value_dict = {
+                                        "type_id": str(type_id),
+                                        "input_type": input_type,
+                                        "output_type": output_type,
+                                        "hasbias": hasbias,
+                                        "Activation": act_tag[1],
+                                        "TileShape": tile_shape[0],
+                                        "ClusterShape": tile_shape[1],
+                                        "KernelSchedule": kernel_schedule,
+                                        "EpilogueSchedule": epilogue_schedule,
+                                        "TileSchedule": tile_schedule,
+                                        "SM": sm,
+                                        "sm": sm[-2:],
+                                    }
+                                    source_all_code += SubstituteTemplate(LaunchGemmPart1, value_dict)
+                                    type_id += 1
+                    source_all_code += LaunchGemmPart2
+                    gemm_config_str = gemm_config_str.replace("<", "").replace(">", "")
+                    code_map[gemm_config_str] = source_all_code
+                    source_path = os.path.join(
+                        generate_dir, f"launch_block_gemm_kernel_sm{sm[-2:]}_{gemm_config_str}.cu"
+                    )
+                    with open(source_path, "w") as f:
+                        f.write(source_all_code)
+                        f.close()
+
+    return head_all_code, code_map
+
+
+# generate fp8_fp8_gemm_scale_bias_act_sm90.cu
+def generate_dispatch_gemm_cu(inputs_type: (str), outputs_type: (str), fuse_gemm_configs: tuple, sm: str):
+    act_tags = [single_config[1] for single_config in fuse_gemm_configs]
+
+    single_config = fuse_gemm_configs[0]
+    hasbiases: (str) = single_config[0]
+    tiles: (str) = single_config[2]
+    KernelSchedule: (str) = single_config[3]
+    EpilogueSchedule: (str) = single_config[4]
+    TileSchedule: (str) = single_config[5]
+    all_code = SubstituteTemplate(code_part0, {"sm": sm[-2:]})
+    type_id = 0
+    for input_type in inputs_type:
+        for output_type in outputs_type:
+            for act_tag in act_tags:
+                for hasbias in hasbiases:
+                    value_dict = {
+                        "act_tag": act_tag[0],
+                        "input_type": input_type,
+                        "output_type": output_type,
+                        "hasbias": hasbias,
+                        "type_id": str(type_id),
+                    }
+                    all_code += SubstituteTemplate(code_part1, value_dict)
+                    type_id += 1
+
+    all_code += code_part2
+    tile_id = 0
+    for tile_shape in tiles:
+        for kernel_schedule in KernelSchedule:
+            for epilogue_schedule in EpilogueSchedule:
+                for tile_schedule in TileSchedule:
+                    if not check_config_valid(tile_shape, kernel_schedule, epilogue_schedule, tile_schedule):
+                        continue
+                    value_dict = {
+                        "TileShape": tile_shape[0],
+                        "ClusterShape": tile_shape[1],
+                        "kernel_schedule": kernel_schedule,
+                        "epilogue_schedule": epilogue_schedule,
+                        "tile_schedule": tile_schedule,
+                        "tile_id": str(tile_id),
+                    }
+                    all_code += SubstituteTemplate(code_part3, value_dict)
+                    tile_id += 1
+    all_code += SubstituteTemplate(code_part4, {"sm": sm[-2:]})
+    tile_id = 0
+    for tile_shape in tiles:
+        blocks, clusters = get_shape_str(tile_shape)
+        gemm_config_str_0 = f"tile{blocks[0]}x{blocks[1]}x{blocks[2]}_cluster{clusters[0]}x{clusters[1]}x{clusters[2]}"
+        for kernel_schedule in KernelSchedule:
+            gemm_config_str_1 = gemm_config_str_0 + f"_{kernel_schedule}"
+            for epilogue_schedule in EpilogueSchedule:
+                gemm_config_str_2 = gemm_config_str_1 + f"_{epilogue_schedule}"
+                for tile_schedule in TileSchedule:
+                    if not check_config_valid(tile_shape, kernel_schedule, epilogue_schedule, tile_schedule):
+                        continue
+                    gemm_config_str = gemm_config_str_2 + f"_{tile_schedule}"
+                    value_dict = {
+                        "sm": sm[-2:],
+                        "tile_id": str(tile_id),
+                        "gemm_config": gemm_config_str.replace("<", "").replace(">", ""),
+                    }
+                    all_code += SubstituteTemplate(code_part5, value_dict)
+                    tile_id += 1
+
+    all_code += SubstituteTemplate(code_part6, {"sm": sm[-2:]})
+    return all_code
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    archs = args.cuda_arch
+    inputs_type = (
+        "float8_e4m3fn",
+        # "float8_e5m2",
+    )
+    outputs_type = ("float16", "bfloat16")
+    sm_dict = {"90": "cutlass::arch::Sm90"}
+
+    for sm in archs:
+        if sm == "90":
+            fuse_gemm_configs = get_candidate_configs(sm)
+            for fuse_gemm_config in fuse_gemm_configs:
+                file_name = f"gpu/cutlass_kernels/fp8_gemm_fused/autogen/generic_block_gemm_kernel_sm{sm}_{fuse_gemm_config[1][0]}.cu"
+                all_code = generate_source_cu(
+                    inputs_type,
+                    outputs_type,
+                    fuse_gemm_config[0],
+                    fuse_gemm_config[1],
+                    fuse_gemm_config[2],
+                    fuse_gemm_config[3],
+                    fuse_gemm_config[4],
+                    fuse_gemm_config[5],
+                    sm_dict[sm],
+                )
+                file_dir = os.path.dirname(file_name)
+                os.makedirs(file_dir, exist_ok=True)
+                with open(file_name, "w") as f:
+                    f.write(all_code)
+                    f.close()
+            # Compile parallelization
+            generate_launch_gemm_cus(
+                "gpu/cutlass_kernels/fp8_gemm_fused/autogen", inputs_type, outputs_type, fuse_gemm_configs, sm_dict[sm]
+            )
+
+            # hard code for act_tag
+            file_name = f"gpu/cutlass_kernels/fp8_gemm_fused/autogen/fp8_fp8_block_gemm_scale_bias_act_sm{sm}.cu"
+            all_code = generate_dispatch_gemm_cu(
+                inputs_type,
+                outputs_type,
+                fuse_gemm_configs,
+                sm_dict[sm],
+            )
+            file_dir = os.path.dirname(file_name)
+            os.makedirs(file_dir, exist_ok=True)
+            with open(file_name, "w") as f:
+                f.write(all_code)
+                f.close()
+        else:
+            raise ValueError(f"Unsupported SM: {sm}")

csrc/utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels.py

@@ -215,10 +215,10 @@ bool fp8_fp8_dual_gemm_scale_bias_act(DualGemmEpilogueAllParams params) {
   int N = params.N;
   int K = params.K;
 
-  std::string mnk_string = "dual_gemm<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
-  std::string regex_mnk_string = "dual_gemm<(\\d+), " + std::to_string(N) + ", " + std::to_string(K) + ">";
-  std::string mnk_split_k_string =  "dual_gemm<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">" + ", split_k";
-  std::string regex_mnk_split_k_string = "dual_gemm<(\\d+), " + std::to_string(N) + ", " + std::to_string(K) + ">, split_k";
+  std::string mnk_string = "tensor_dual_gemm_sm89<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
+  std::string regex_mnk_string = "tensor_dual_gemm_sm89<(\\d+), " + std::to_string(N) + ", " + std::to_string(K) + ">";
+  std::string mnk_split_k_string =  "tensor_dual_gemm_sm89<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">" + ", split_k";
+  std::string regex_mnk_split_k_string = "tensor_dual_gemm_sm89<(\\d+), " + std::to_string(N) + ", " + std::to_string(K) + ">, split_k";
   int split_k;
   int kernel_id;
   std::string best_config;

csrc/utils/auto_gen_fp8_fp8_dual_gemm_fused_kernels_sm90.py

@@ -211,7 +211,7 @@ bool fp8_fp8_dual_gemm_scale_bias_act(DualGemmEpilogueAllParams params) {
   int K = params.K;
 
   int kernel_id;
-  std::string mnk_string = "dual_gemm<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
+  std::string mnk_string = "tensor_dual_gemm_sm90<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
   std::string best_config;
   CutlassGemmConfigMannager& best_config_mannager = CutlassGemmConfigMannager::getInstance();
   char *config_file_path_c_str = getenv("FLAGS_use_cutlass_device_best_config_path");

csrc/utils/auto_gen_fp8_fp8_gemm_fused_kernels.py

@@ -232,10 +232,10 @@ bool fp8_fp8_gemm_scale_bias_act(GemmEpilogueAllParams params) {
   int N = params.N;
   int K = params.K;
 
-  std::string mnk_string = "gemm<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
-  std::string regex_mnk_string = "gemm<(\\d+), " + std::to_string(N) + ", " + std::to_string(K) + ">";
-  std::string mnk_split_k_string =  "gemm<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">" + ", split_k";
-  std::string regex_mnk_split_k_string = "gemm<(\\d+), " + std::to_string(N) + ", " + std::to_string(K) + ">, split_k";
+  std::string mnk_string = "tensor_gemm_sm90<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
+  std::string regex_mnk_string = "tensor_gemm_sm90<(\\d+), " + std::to_string(N) + ", " + std::to_string(K) + ">";
+  std::string mnk_split_k_string =  "tensor_gemm_sm90<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">" + ", split_k";
+  std::string regex_mnk_split_k_string = "tensor_gemm_sm90<(\\d+), " + std::to_string(N) + ", " + std::to_string(K) + ">, split_k";
   int split_k;
   int kernel_id;
   std::string best_config;

csrc/utils/auto_gen_fp8_fp8_gemm_fused_kernels_ptr_scale_sm90.py

@@ -0,0 +1,582 @@
+# 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 argparse
+import copy
+import os
+import re
+
+
+def get_candidate_tiles():
+    base_configs = [
+        ("<_64, _64, _128>", "<_1, _8, _1>"),
+        ("<_64, _128, _128>", "<_2, _1, _1>"),
+        ("<_128, _128, _128>", "<_2, _1, _1>"),
+    ]
+    base_configs.extend(
+        [
+            ("<_64, _64, _128>", "<_1, _1, _1>"),
+            ("<_64, _64, _128>", "<_1, _2, _1>"),
+            ("<_64, _64, _128>", "<_2, _1, _1>"),
+            ("<_64, _64, _64>", "<_1, _1, _1>"),
+            ("<_64, _64, _64>", "<_1, _2, _1>"),
+            ("<_64, _64, _64>", "<_2, _1, _1>"),
+            ("<_64, _128, _128>", "<_1, _2, _1>"),
+            ("<_64, _128, _128>", "<_1, _1, _1>"),
+            ("<_128, _128, _64>", "<_2, _1, _1>"),
+            ("<_256, _128, _128>", "<_1, _2, _1>"),
+            ("<_256, _128, _128>", "<_1, _1, _1>"),
+            # The following configurations are rarely selected in Qwen2-7B-model.
+            # ("<_256, _128, _128>", "<_4, _1, _1>"),
+            # ("<_256, _128, _128>", "<_1, _4, _1>"),
+            # ("<_256, _128, _128>", "<_2, _4, _1>"),
+            # ("<_128, _128, _256>", "<_1, _2, _1>"),
+            # ("<_128, _128, _128>", "<_4, _1, _1>"),
+            # ("<_128, _128, _128>", "<_2, _4, _1>"),
+            # ("<_128, _128, _128>", "<_1, _2, _1>"),
+            # ("<_128, _128, _128>", "<_1, _1, _1>"),
+            # ("<_128, _128, _128>", "<_1, _4, _1>"),
+            # ("<_128, _128, _64>", "<_2, _2, _1>"),
+        ]
+    )
+
+    return base_configs
+
+
+def get_candidate_configs(sm):
+    tiles = get_candidate_tiles()
+    candidate_configs = list()
+
+    hasbias = ("false", "true")
+    KernelSchedule = (
+        "KernelTmaWarpSpecializedFP8FastAccum",
+        "KernelTmaWarpSpecializedPingpongFP8FastAccum",
+        # "KernelTmaWarpSpecializedCooperativeFP8FastAccum",
+    )
+    EpilogueSchedule = ("TmaWarpSpecialized", "TmaWarpSpecializedCooperative")
+    for act_tag in [
+        ("noact", "Identity"),
+        # ("relu", "ReLu"),
+        # ("gelu", "GELU"),
+    ]:
+        candidate_configs.extend([(hasbias, act_tag, tiles, KernelSchedule, EpilogueSchedule)])
+
+    return candidate_configs
+
+
+def get_shape_str(tile_shape):
+    blocks, clusters = [s.replace(" ", "").strip("<>").split(",") for s in tile_shape]
+    blocks = [elem.strip("_") for elem in blocks]
+    clusters = [elem.strip("_") for elem in clusters]
+    return blocks, clusters
+
+
+def check_config_valid(tile_shape, kernel_schedule, epilogue_schedule):
+    blocks, clusters = get_shape_str(tile_shape)
+    if int(blocks[0]) < 128 and kernel_schedule == "KernelTmaWarpSpecializedCooperativeFP8FastAccum":
+        return False
+    if "Cooperative" in kernel_schedule and "Cooperative" not in epilogue_schedule:
+        return False
+    if (
+        tile_shape[0] == "<_256, _128, _128>"
+        and "Cooperative" not in kernel_schedule
+        and "Cooperative" not in epilogue_schedule
+    ):
+        return False
+    # flag1 = (int(blocks[0]) == 64 and kernel_schedule == "KernelTmaWarpSpecializedPingpongFP8FastAccum" and epilogue_schedule == "TmaWarpSpecialized")
+    # flag2 = (int(blocks[0]) != 64 and kernel_schedule == "KernelTmaWarpSpecializedCooperativeFP8FastAccum" and epilogue_schedule == "TmaWarpSpecializedCooperative")
+    # if not (flag1 or flag2):
+    #     return False
+    return True
+
+
+# this is a file's header part
+CommonHead = """// Generated by auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py - Do not edit.
+
+#pragma once
+
+#include "fp8_gemm_fused/fuse_gemm_act_template_3x.h"
+
+"""
+
+GemmDeclare = """
+template<>
+bool dispatch_fp8_gemm_ptr_scale_sm{sm}<phi::dtype::{input_type}, phi::dtype::{output_type},
+                                 {hasbias},
+                                 Shape{TileShape},
+                                 Shape{ClusterShape},
+                                 cutlass::gemm::{KernelSchedule},
+                                 cutlass::epilogue::{EpilogueSchedule},
+                                 {SM}
+                                 >(GemmEpilogueAllParams);
+
+"""
+
+
+LaunchGemmHead = """
+#pragma once
+
+#include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
+
+"""
+
+LaunchGemmDeclare = """
+bool launch_gemm_kernel_ptr_scale_sm{sm}_{gemm_config}(const int type_id, GemmEpilogueAllParams params);
+"""
+
+LaunchGemmPart0 = """
+#pragma once
+
+#include "launch_gemm_kernel_ptr_scale_sm{sm}.h"
+
+bool launch_gemm_kernel_ptr_scale_sm{sm}_{gemm_config}(const int type_id, GemmEpilogueAllParams params){
+    switch (type_id) {
+"""
+
+LaunchGemmPart1 = """
+        case {type_id}:
+            return dispatch_fp8_gemm_ptr_scale_sm{sm}<phi::dtype::{input_type}, phi::dtype::{output_type},
+                                 {hasbias},
+                                 Shape{TileShape},
+                                 Shape{ClusterShape},
+                                 cutlass::gemm::{KernelSchedule},
+                                 cutlass::epilogue::{EpilogueSchedule},
+                                 {SM}
+                                 >(params);
+            break;
+"""
+
+LaunchGemmPart2 = """
+        default:
+            throw std::runtime_error("cutlass gemm config is invalid.");
+            break;
+    }
+    return false;
+}
+"""
+
+
+code_part0 = """// Generated by auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py - Do not edit.
+
+#include <map>
+#include "fp8_gemm_fused/fp8_fp8_gemm_scale_bias_act.h"
+#include "launch_gemm_kernel_ptr_scale_sm{sm}.h"
+
+COMMON_DECLARE_string(use_cutlass_device_best_config_path);
+
+std::map<std::string, int> gemm_ptr_scale_type_map{"""
+
+code_part1 = """
+    {"{input_type}_{output_type}_{hasbias}_{act_tag}",   {type_id}}, """
+
+code_part2 = """
+};
+
+std::map<std::string, int> gemm_ptr_scale_config_map{
+"""
+
+code_part3 = """    {"{TileShape}, {ClusterShape}, {kernel_schedule}, {epilogue_schedule}", {tile_id}},
+"""
+
+code_part4 = """};
+
+bool launch_gemm_kernel_ptr_scale_sm{sm}(const int type_id, const int kernel_id, GemmEpilogueAllParams params){
+    switch (kernel_id) {"""
+
+code_part5 = """
+        case {tile_id}:
+            return launch_gemm_kernel_ptr_scale_sm{sm}_{gemm_config}(type_id, params);
+            break;
+"""
+
+code_part6 = """
+        default:
+            throw std::runtime_error("fp8 gemm_fused Config is invalid.");
+            break;
+    }
+    return false;
+}
+
+template <typename T>
+T get_relative_best(nlohmann::json* json_data,
+                    const std::string& target_key,
+                    const int& m) {
+    if (json_data->contains(target_key)) {
+        return json_data->at(target_key);
+    } else {
+        if (m <= 64){
+            return "<_64, _64, _128>, <_1, _8, _1>, KernelTmaWarpSpecializedPingpongFP8FastAccum, TmaWarpSpecialized";
+        }else if(m <= 128){
+            return "<_64, _128, _128>, <_2, _1, _1>, KernelTmaWarpSpecializedPingpongFP8FastAccum, TmaWarpSpecialized";
+        }else{
+            return "<_128, _128, _128>, <_2, _1, _1>, KernelTmaWarpSpecializedPingpongFP8FastAccum, TmaWarpSpecialized";
+        }
+    }
+}
+
+bool fp8_fp8_gemm_ptr_scale_bias_act(GemmEpilogueAllParams params) {
+  if (gemm_ptr_scale_type_map.find(params.fuse_gemm_config) == gemm_ptr_scale_type_map.end()) {
+    throw std::runtime_error("fp8 gemm_fused_ptr_scale config is invalid.");
+  }
+
+  int type_id = gemm_ptr_scale_type_map[params.fuse_gemm_config];
+  int M = (params.M + 31) / 32 * 32;
+  int N = params.N;
+  int K = params.K;
+
+  int kernel_id;
+  std::string mnk_string = "tensor_gemm_ptr_scale_sm90<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
+  std::string best_config;
+  CutlassGemmConfigMannager& best_config_mannager = CutlassGemmConfigMannager::getInstance();
+  char *config_file_path_c_str = getenv("FLAGS_use_cutlass_device_best_config_path");
+  std::string config_file_path = config_file_path_c_str == nullptr ? "" : std::string(config_file_path_c_str);
+  if(config_file_path == "tune"){ // tune kernel
+    int warm_up_times = 5;
+    int tune_times = 10;
+    std::string best_kernel_id = "";
+    float duratation = 1000000.f;
+    // tune all kernel_id kernels
+    for(const auto& config_pair : gemm_ptr_scale_config_map){
+        std::cout << "Running tune kernel: " << config_pair.first<< std::endl;
+        bool is_valid = true;
+        // warm up
+        for(int num_time = 0; num_time < warm_up_times; ++num_time){
+            if(!launch_gemm_kernel_ptr_scale_sm{sm}(type_id, config_pair.second, params)){
+                is_valid = false;
+                break;
+            }
+        }
+        if(!is_valid){
+            continue;
+        }
+        cudaEvent_t start, stop;
+        cudaEventCreate(&start);
+        cudaEventCreate(&stop);
+        cudaStreamSynchronize(params.stream);
+        cudaEventRecord(start, params.stream);
+        for(int num_time = 0; num_time < tune_times; ++num_time){
+            if(!launch_gemm_kernel_ptr_scale_sm{sm}(type_id, config_pair.second, params)){
+                is_valid = false;
+                break;
+            };
+        }
+        cudaEventRecord(stop, params.stream);
+        cudaEventSynchronize(stop);
+        float elapsedTime;
+        if(is_valid){
+            cudaEventElapsedTime(&elapsedTime, start, stop);
+        } else {
+            continue;
+        }
+        cudaEventDestroy(start);
+        cudaEventDestroy(stop);
+        if(elapsedTime < duratation){
+            best_kernel_id = config_pair.first;
+            duratation = elapsedTime;
+        }
+    }
+
+    nlohmann::json new_json;
+    new_json[mnk_string] = best_kernel_id;
+    best_config_mannager.up_date_configs(new_json);
+    std::cout <<"Gemm tune result for " << mnk_string<< ": best config is: "<< best_kernel_id << std::endl;
+    return true;
+  } else { // run kernel
+    nlohmann::json* config_json = new nlohmann::json();
+    if (config_file_path != "" && config_file_path != "default") {
+        config_json = best_config_mannager.get_gemm_best_configs(config_file_path);
+    }
+
+    best_config = get_relative_best<std::string>(config_json, mnk_string, M);
+
+    if (gemm_ptr_scale_config_map.find(best_config) == gemm_ptr_scale_config_map.end()) {
+        throw std::runtime_error("This config'kernel not be generate, please check auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py and re-generate.");
+    } else {
+        kernel_id = gemm_ptr_scale_config_map[best_config];
+    }
+    return launch_gemm_kernel_ptr_scale_sm{sm}(type_id, kernel_id, params);
+  }
+}
+
+"""
+
+
+def SubstituteTemplate(template, values_base):
+    values = copy.deepcopy(values_base)
+    if values.get("KernelSchedule") is not None and "Auto" in values["KernelSchedule"]:
+        values["KernelSchedule"] = "collective::" + values["KernelSchedule"]
+    if values.get("EpilogueSchedule") is not None and "Auto" in values["EpilogueSchedule"]:
+        values["EpilogueSchedule"] = "collective::" + values["EpilogueSchedule"]
+    text = template
+    changed = True
+    while changed:
+        changed = False
+        for key, value in values.items():
+            regex = f"\\{{{key}\\}}"
+            newtext = re.sub(regex, value, text)
+            if newtext != text:
+                changed = True
+            text = newtext
+    return text
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="The argument for generating the generic_mixed_gemm_kernelLauncher instance."
+    )
+    parser.add_argument(
+        "--cuda_arch",
+        type=str,
+        nargs="+",
+        default=["90"],
+        help="The CUDA architecture to be generated.",
+    )
+
+    args = parser.parse_args()
+    return args
+
+
+# generate source .cu
+def generate_source_cu(
+    inputs_type: (str),
+    outputs_type: (str),
+    hasbiases: (str),
+    act_tag: (str),
+    tiles: (str),
+    KernelSchedule: (str),
+    EpilogueSchedule: (str),
+    sm: str,
+):
+    all_code = CommonHead
+
+    for input_type in inputs_type:
+        for output_type in outputs_type:
+            for hasbias in hasbiases:
+                for tile_config in tiles:
+                    for kernel_schedule in KernelSchedule:
+                        for epilogue_schedule in EpilogueSchedule:
+                            if not check_config_valid(tile_config, kernel_schedule, epilogue_schedule):
+                                continue
+                            value_dict = {
+                                "input_type": input_type,
+                                "output_type": output_type,
+                                "hasbias": hasbias,
+                                "Activation": act_tag[1],
+                                "TileShape": tile_config[0],
+                                "ClusterShape": tile_config[1],
+                                "KernelSchedule": kernel_schedule,
+                                "EpilogueSchedule": epilogue_schedule,
+                                "SM": sm,
+                                "sm": sm[-2:],
+                            }
+                            all_code += SubstituteTemplate(GemmDeclare, value_dict)
+
+    return all_code
+
+
+# generate gemm launch .cu
+def generate_launch_gemm_cus(
+    generate_dir: (str), inputs_type: (str), outputs_type: (str), fuse_gemm_configs: tuple, sm: str
+):
+    act_tags = [single_config[1] for single_config in fuse_gemm_configs]
+
+    single_config = fuse_gemm_configs[0]
+    hasbiases: (str) = single_config[0]
+    tiles: (str) = single_config[2]
+    KernelSchedule: (str) = single_config[3]
+    EpilogueSchedule: (str) = single_config[4]
+    code_map = {}
+    head_path = os.path.join(generate_dir, f"launch_gemm_kernel_ptr_scale_sm{sm[-2:]}.h")
+    head_all_code = LaunchGemmHead
+    for tile_config in tiles:
+        blocks, clusters = get_shape_str(tile_config)
+        gemm_config_str_0 = f"tile{blocks[0]}x{blocks[1]}x{blocks[2]}_cluster{clusters[0]}x{clusters[1]}x{clusters[2]}"
+        for kernel_schedule in KernelSchedule:
+            gemm_config_str_1 = gemm_config_str_0 + f"_{kernel_schedule}"
+            for epilogue_schedule in EpilogueSchedule:
+                if not check_config_valid(tile_config, kernel_schedule, epilogue_schedule):
+                    continue
+                gemm_config_str = gemm_config_str_1 + f"_{epilogue_schedule}"
+                value_dict = {
+                    "sm": sm[-2:],
+                    "gemm_config": gemm_config_str,
+                }
+                head_all_code += SubstituteTemplate(LaunchGemmDeclare, value_dict)
+    os.makedirs(generate_dir, exist_ok=True)
+    with open(head_path, "w") as f:
+        f.write(head_all_code)
+        f.close()
+
+    for tile_shape in tiles:
+        blocks, clusters = get_shape_str(tile_shape)
+        gemm_config_str_0 = f"tile{blocks[0]}x{blocks[1]}x{blocks[2]}_cluster{clusters[0]}x{clusters[1]}x{clusters[2]}"
+        for kernel_schedule in KernelSchedule:
+            gemm_config_str_1 = gemm_config_str_0 + f"_{kernel_schedule}"
+            for epilogue_schedule in EpilogueSchedule:
+                if not check_config_valid(tile_shape, kernel_schedule, epilogue_schedule):
+                    continue
+                gemm_config_str = gemm_config_str_1 + f"_{epilogue_schedule}"
+                value_dict = {
+                    "sm": sm[-2:],
+                    "gemm_config": gemm_config_str,
+                }
+                source_all_code = SubstituteTemplate(LaunchGemmPart0, value_dict)
+                type_id = 0
+                for input_type in inputs_type:
+                    for output_type in outputs_type:
+                        for act_tag in act_tags:
+                            for hasbias in hasbiases:
+                                value_dict = {
+                                    "type_id": str(type_id),
+                                    "input_type": input_type,
+                                    "output_type": output_type,
+                                    "hasbias": hasbias,
+                                    "Activation": act_tag[1],
+                                    "TileShape": tile_shape[0],
+                                    "ClusterShape": tile_shape[1],
+                                    "KernelSchedule": kernel_schedule,
+                                    "EpilogueSchedule": epilogue_schedule,
+                                    "SM": sm,
+                                    "sm": sm[-2:],
+                                }
+                                source_all_code += SubstituteTemplate(LaunchGemmPart1, value_dict)
+                                type_id += 1
+                source_all_code += LaunchGemmPart2
+                code_map[gemm_config_str] = source_all_code
+                source_path = os.path.join(
+                    generate_dir, f"launch_gemm_kernel_ptr_scale_sm{sm[-2:]}_{gemm_config_str}.cu"
+                )
+                with open(source_path, "w") as f:
+                    f.write(source_all_code)
+                    f.close()
+
+    return head_all_code, code_map
+
+
+# generate fp8_fp8_gemm_ptr_scale_bias_act_sm90.cu
+def generate_dispatch_gemm_cu(inputs_type: (str), outputs_type: (str), fuse_gemm_configs: tuple, sm: str):
+    act_tags = [single_config[1] for single_config in fuse_gemm_configs]
+
+    single_config = fuse_gemm_configs[0]
+    hasbiases: (str) = single_config[0]
+    tiles: (str) = single_config[2]
+    KernelSchedule: (str) = single_config[3]
+    EpilogueSchedule: (str) = single_config[4]
+
+    all_code = SubstituteTemplate(code_part0, {"sm": sm[-2:]})
+    type_id = 0
+    for input_type in inputs_type:
+        for output_type in outputs_type:
+            for act_tag in act_tags:
+                for hasbias in hasbiases:
+                    value_dict = {
+                        "act_tag": act_tag[0],
+                        "input_type": input_type,
+                        "output_type": output_type,
+                        "hasbias": hasbias,
+                        "type_id": str(type_id),
+                    }
+                    all_code += SubstituteTemplate(code_part1, value_dict)
+                    type_id += 1
+
+    all_code += code_part2
+    tile_id = 0
+    for tile_shape in tiles:
+        for kernel_schedule in KernelSchedule:
+            for epilogue_schedule in EpilogueSchedule:
+                if not check_config_valid(tile_shape, kernel_schedule, epilogue_schedule):
+                    continue
+                value_dict = {
+                    "TileShape": tile_shape[0],
+                    "ClusterShape": tile_shape[1],
+                    "kernel_schedule": kernel_schedule,
+                    "epilogue_schedule": epilogue_schedule,
+                    "tile_id": str(tile_id),
+                }
+                all_code += SubstituteTemplate(code_part3, value_dict)
+                tile_id += 1
+    all_code += SubstituteTemplate(code_part4, {"sm": sm[-2:]})
+    tile_id = 0
+    for tile_shape in tiles:
+        blocks, clusters = get_shape_str(tile_shape)
+        gemm_config_str_0 = f"tile{blocks[0]}x{blocks[1]}x{blocks[2]}_cluster{clusters[0]}x{clusters[1]}x{clusters[2]}"
+        for kernel_schedule in KernelSchedule:
+            gemm_config_str_1 = gemm_config_str_0 + f"_{kernel_schedule}"
+            for epilogue_schedule in EpilogueSchedule:
+                if not check_config_valid(tile_shape, kernel_schedule, epilogue_schedule):
+                    continue
+                gemm_config_str = gemm_config_str_1 + f"_{epilogue_schedule}"
+                value_dict = {
+                    "sm": sm[-2:],
+                    "tile_id": str(tile_id),
+                    "gemm_config": gemm_config_str,
+                }
+                all_code += SubstituteTemplate(code_part5, value_dict)
+                tile_id += 1
+
+    all_code += SubstituteTemplate(code_part6, {"sm": sm[-2:]})
+    return all_code
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    archs = args.cuda_arch
+    inputs_type = (
+        "float8_e4m3fn",
+        "float8_e5m2",
+    )
+    outputs_type = ("float16", "bfloat16")
+    sm_dict = {"90": "cutlass::arch::Sm90"}
+
+    for sm in archs:
+        if sm == "90":
+            fuse_gemm_configs = get_candidate_configs(sm)
+            for fuse_gemm_config in fuse_gemm_configs:
+                file_name = f"gpu/cutlass_kernels/fp8_gemm_fused/autogen/generic_gemm_kernel_ptr_scale_sm{sm}_{fuse_gemm_config[1][0]}.cu"
+                all_code = generate_source_cu(
+                    inputs_type,
+                    outputs_type,
+                    fuse_gemm_config[0],
+                    fuse_gemm_config[1],
+                    fuse_gemm_config[2],
+                    fuse_gemm_config[3],
+                    fuse_gemm_config[4],
+                    sm_dict[sm],
+                )
+                file_dir = os.path.dirname(file_name)
+                os.makedirs(file_dir, exist_ok=True)
+                with open(file_name, "w") as f:
+                    f.write(all_code)
+                    f.close()
+            # Compile parallelization
+            generate_launch_gemm_cus(
+                "gpu/cutlass_kernels/fp8_gemm_fused/autogen", inputs_type, outputs_type, fuse_gemm_configs, sm_dict[sm]
+            )
+
+            # hard code for act_tag
+            file_name = f"gpu/cutlass_kernels/fp8_gemm_fused/autogen/fp8_fp8_gemm_ptr_scale_bias_act_sm{sm}.cu"
+            all_code = generate_dispatch_gemm_cu(
+                inputs_type,
+                outputs_type,
+                fuse_gemm_configs,
+                sm_dict[sm],
+            )
+            file_dir = os.path.dirname(file_name)
+            os.makedirs(file_dir, exist_ok=True)
+            with open(file_name, "w") as f:
+                f.write(all_code)
+                f.close()
+        else:
+            raise ValueError(f"Unsupported SM: {sm}")

csrc/utils/auto_gen_fp8_fp8_gemm_fused_kernels_sm90.py

@@ -239,7 +239,7 @@ bool fp8_fp8_gemm_scale_bias_act(GemmEpilogueAllParams params) {
   int K = params.K;
 
   int kernel_id;
-  std::string mnk_string = "gemm<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
+  std::string mnk_string = "tensor_gemm_sm90<"+ std::to_string(M)+ ", " +std::to_string(N) + ", "+ std::to_string(K)+ ">";
   std::string best_config;
   CutlassGemmConfigMannager& best_config_mannager = CutlassGemmConfigMannager::getInstance();
   char *config_file_path_c_str = getenv("FLAGS_use_cutlass_device_best_config_path");

csrc/utils/tune_cutlass_fp8_block_gemm.py

@@ -0,0 +1,60 @@
+# 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 argparse
+
+import paddle
+from paddlenlp_ops import cutlass_fp8_fp8_half_block_gemm_fused
+
+
+def setup_args():
+    """Setup export arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--m_min", type=int, help="range of gemm shape: m_min")
+    parser.add_argument("--m_max", type=int, help="range of gemm shape: m_max")
+    parser.add_argument("--n", nargs="+", type=int, help="List of gemm shape: n")
+    parser.add_argument("--k", nargs="+", type=int, help="List of gemm shape: k")
+    args = parser.parse_args()
+    return args
+
+
+def gemm(m, n, k):
+    scale_n = (n + 128 - 1) // 128
+    scale_k = (k + 128 - 1) // 128
+    A = paddle.ones([m, k], dtype="float8_e4m3fn")
+    B = paddle.ones([n, k], dtype="float8_e4m3fn")
+    a_scale = paddle.randn([scale_k, m], dtype="float32")
+    b_scale = paddle.randn([scale_n, scale_k], dtype="float32")
+    res = cutlass_fp8_fp8_half_block_gemm_fused(
+        A, B, a_scale, b_scale, bias=None, transpose_x=False, transpose_y=True, output_dtype="bfloat16", act="identity"
+    )
+    return res
+
+
+if __name__ == "__main__":
+    args = setup_args()
+
+    m_min = args.m_min
+    m_max = args.m_max
+    ns = args.n
+    ks = args.k
+
+    for m in range(m_min, m_max + 32, 32):
+        if m > m_max:
+            break
+        for idx in range(len(ns)):
+            n = ns[idx]
+            k = ks[idx]
+            gemm(m, n, k)
+            paddle.device.cuda.empty_cache()

csrc/utils/tune_cutlass_fp8_gemm_ptr_scale.py

@@ -0,0 +1,65 @@
+# 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 argparse
+
+import paddle
+from paddlenlp_ops import cutlass_fp8_fp8_half_gemm_ptr_scale_fused
+
+
+def setup_args():
+    """Setup export arguments."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--m_min", type=int, help="range of gemm shape: m_min")
+    parser.add_argument("--m_max", type=int, help="range of gemm shape: m_max")
+    parser.add_argument("--n", nargs="+", type=int, help="List of gemm shape: n")
+    parser.add_argument("--k", nargs="+", type=int, help="List of gemm shape: k")
+    args = parser.parse_args()
+    return args
+
+
+def gemm(m, n, k):
+    A = paddle.ones([m, k], dtype="float8_e4m3fn")
+    B = paddle.ones([n, k], dtype="float8_e4m3fn")
+    A_s = paddle.rand([1], dtype="float32")
+    B_s = paddle.rand([1], dtype="float32")
+    res = cutlass_fp8_fp8_half_gemm_ptr_scale_fused(
+        A,
+        B,
+        x_scale=A_s,
+        y_scale=B_s,
+        bias=None,
+        transpose_x=False,
+        transpose_y=True,
+        output_dtype="bfloat16",
+    )
+    return res
+
+
+if __name__ == "__main__":
+    args = setup_args()
+
+    m_min = args.m_min
+    m_max = args.m_max
+    ns = args.n
+    ks = args.k
+
+    for m in range(m_min, m_max + 32, 32):
+        if m > m_max:
+            break
+        for idx in range(len(ns)):
+            n = ns[idx]
+            k = ks[idx]
+            gemm(m, n, k)
+            paddle.device.cuda.empty_cache()