Performance Update (2025.04.22) (#71)

* Fix benchmark script

* Performance optimization for compute-bound cases

* Add new testcase (s_k = 16384)

* Update README.md

* Update comment

* Update README.md

* Add the deep-dive blog

* Add background color for MLA Kernel Sched.drawio.svg

* Use relative path for the schedule image

* Move flash_mla.h to kernels/params.h

.gitignore

@@ -5,3 +5,4 @@ __pycache__/
 dist/
 *perf.csv
 *.png
+/.vscode

README.md

 # FlashMLA
 
+## Performance Update (2025.04.22)
+
+We're excited to announce the new release of Flash MLA, which delivers 5% ~ 15% performance improvement on compute-bound workloads, achieving up to 660 TFlops on NVIDIA H800 SXM5 GPUs. The interface of the new version is fully compatible with the old one. Just switch to the new version and enjoy the instant speedup! 🚀🚀🚀
+
+Besides, we'd love to share the technical details behind the new kernel! Check out our deep-dive write-up here: <LINK>
+
+The new kernel primarily targets compute-intensive settings (where the number of q heads $\times$ the number of q tokens per request (if MTP is disabled then it's 1) $\ge 64$). For memory-bound cases, we recommend using version [b31bfe7](https://github.com/deepseek-ai/FlashMLA/tree/b31bfe72a83ea205467b3271a5845440a03ed7cb) for optimal performance.
+
+## Introduction
+
 FlashMLA is an efficient MLA decoding kernel for Hopper GPUs, optimized for variable-length sequences serving.
 
 Currently released:
 - BF16, FP16
 - Paged kvcache with block size of 64
 
+## Requirements
+
+- Hopper GPUs
+- CUDA 12.3 and above
+    - **But we highly recommend 12.8 or above for the best performance**
+- PyTorch 2.0 and above
+
 ## Quick start
 
 ### Install
@@ -20,7 +37,9 @@ python setup.py install
 python tests/test_flash_mla.py
 ```
 
-Achieving up to 3000 GB/s in memory-bound configuration and 580 TFLOPS in computation-bound configuration on H800 SXM5, using CUDA 12.8.
+It is able up to 3000 GB/s in memory-bound configuration and 660 TFLOPS in computation-bound configuration on H800 SXM5, using CUDA 12.8.
+
+Note. For memory-bound cases, we recommend using version [b31bfe7](https://github.com/deepseek-ai/FlashMLA/tree/b31bfe72a83ea205467b3271a5845440a03ed7cb) for optimal performance.
 
 ### Usage
 
@@ -38,13 +57,6 @@ for i in range(num_layers):
     ...
 ```
 
-## Requirements
-
-- Hopper GPUs
-- CUDA 12.3 and above
-    - **But we highly recommend 12.8 or above for the best performance**
-- PyTorch 2.0 and above
-
 ## Acknowledgement
 
 FlashMLA is inspired by [FlashAttention 2&3](https://github.com/dao-AILab/flash-attention/) and [cutlass](https://github.com/nvidia/cutlass) projects.
@@ -91,7 +103,7 @@ The corresponding FlashMLA version can be found at: [AITER/MLA](https://github.c
 ```bibtex
 @misc{flashmla2025,
       title={FlashMLA: Efficient MLA decoding kernels},
-      author={Jiashi Li},
+      author={Jiashi Li, Shengyu Liu},
       year={2025},
       publisher = {GitHub},
       howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}},

benchmark/bench_flash_mla.py

@@ -435,7 +435,7 @@ def compare_ab(baseline, target, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal
     out_b, lse_b, perf_b = target_func(q, block_table, blocked_k, max_seqlen_pad, block_size, b, s_q, cache_seqlens, h_q, h_kv, d, dv, causal, dtype)
     
     torch.testing.assert_close(out_b.float(), out_a.float(), atol=1e-2, rtol=1e-2), "out"
-    if target not in ["flash_infer", "flash_mla_triton"]:
+    if target not in ["flash_infer", "flash_mla_triton"] and baseline not in ["flash_infer", "flash_mla_triton"]:
         # flash_infer has a different lse return value
         # flash_mla_triton doesn't return lse
         torch.testing.assert_close(lse_b.float(), lse_a.float(), atol=1e-2, rtol=1e-2), "lse"

csrc/flash_api.cpp

@@ -10,8 +10,11 @@
 
 #include <cutlass/fast_math.h>
 
-#include "flash_mla.h"
-#include "static_switch.h"
+#include "kernels/config.h"
+#include "kernels/get_mla_metadata.h"
+#include "kernels/mla_combine.h"
+#include "kernels/params.h"
+#include "kernels/splitkv_mla.h"
 
 #define CHECK_DEVICE(x) TORCH_CHECK(x.is_cuda(), #x " must be on CUDA")
 #define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
@@ -23,11 +26,6 @@ get_mla_metadata(
     const int num_heads_per_head_k,
     const int num_heads_k
 ) {
-    // This should match the logic in the MLA kernel.
-    static constexpr int block_size_m = 64;
-    static constexpr int block_size_n = 64;
-    static constexpr int fixed_overhead_num_blocks = 5;
-
     CHECK_DEVICE(seqlens_k);
     TORCH_CHECK(seqlens_k.is_contiguous());
     TORCH_CHECK(seqlens_k.dtype() == torch::kInt32);
@@ -38,7 +36,7 @@ get_mla_metadata(
 
     auto dprops = at::cuda::getCurrentDeviceProperties();
     int sm_count = dprops->multiProcessorCount;
-    int num_sm_parts = sm_count / num_heads_k / cutlass::ceil_div(num_heads_per_head_k, block_size_m);
+    int num_sm_parts = sm_count / num_heads_k / cutlass::ceil_div(num_heads_per_head_k, Config::BLOCK_SIZE_M);
 
     auto tile_scheduler_metadata = torch::empty({num_sm_parts, TileSchedulerMetaDataSize}, options);
     auto num_splits = torch::empty({batch_size + 1}, options);
@@ -52,10 +50,10 @@ get_mla_metadata(
     params.tile_scheduler_metadata_ptr = tile_scheduler_metadata_ptr;
     params.num_splits_ptr = num_splits_ptr;
     params.batch_size = batch_size;
-    params.block_size_n = block_size_n;
-    params.fixed_overhead_num_blocks = fixed_overhead_num_blocks;
+    params.block_size_n = Config::PAGE_BLOCK_SIZE;
+    params.fixed_overhead_num_blocks = Config::FIXED_OVERHEAD_NUM_BLOCKS;
     params.num_sm_parts = num_sm_parts;
-    get_mla_metadata_func(params, stream);
+    run_get_mla_metadata_kernel(params, stream);
 
     return {tile_scheduler_metadata, num_splits};
 }
@@ -64,7 +62,6 @@ std::vector<at::Tensor>
 mha_fwd_kvcache_mla(
     at::Tensor &q,                               // batch_size x seqlen_q x num_heads x head_size
     const at::Tensor &kcache,                    // num_blocks x page_block_size x num_heads_k x head_size
-    std::optional<const at::Tensor> &vcache_,    // num_blocks x page_block_size x num_heads_k x head_size_v
     const int head_size_v,
     const at::Tensor &seqlens_k,                 // batch_size
     const at::Tensor &block_table,               // batch_size x max_num_blocks_per_seq
@@ -73,138 +70,141 @@ mha_fwd_kvcache_mla(
     const at::Tensor &tile_scheduler_metadata,   // num_sm_parts x TileSchedulerMetaDataSize
     const at::Tensor &num_splits                 // batch_size + 1
 ) {
+    // Check the architecture
     auto dprops = at::cuda::getCurrentDeviceProperties();
     bool is_sm90 = dprops->major == 9 && dprops->minor == 0;
     TORCH_CHECK(is_sm90);
 
-    at::Tensor vcache = vcache_.has_value() ? vcache_.value() : kcache;
-
+    // Check data types
     auto q_dtype = q.dtype();
+    TORCH_CHECK(q_dtype == torch::kBFloat16 || q_dtype == torch::kHalf);
     TORCH_CHECK(kcache.dtype() == q_dtype, "query and key must have the same dtype");
+    TORCH_CHECK(seqlens_k.dtype() == torch::kInt32, "seqlens_k must have dtype int32");
+    TORCH_CHECK(block_table.dtype() == torch::kInt32, "block_table must have dtype torch.int32");
+    TORCH_CHECK(tile_scheduler_metadata.dtype() == torch::kInt32, "tile_scheduler_metadata must have dtype int32");
+    TORCH_CHECK(num_splits.dtype() == torch::kInt32, "num_splits must have dtype int32");
 
-    CHECK_DEVICE(q); CHECK_DEVICE(kcache); CHECK_DEVICE(vcache);
+    // Check device
+    CHECK_DEVICE(q);
+    CHECK_DEVICE(kcache);
+    CHECK_DEVICE(seqlens_k);
+    CHECK_DEVICE(block_table);
+    CHECK_DEVICE(tile_scheduler_metadata);
+    CHECK_DEVICE(num_splits);
 
+    // Check layout
     TORCH_CHECK(q.stride(-1) == 1, "Input tensor must have contiguous last dimension");
     TORCH_CHECK(kcache.stride(-1) == 1, "Input tensor must have contiguous last dimension");
-    TORCH_CHECK(vcache.stride(-1) == 1, "Input tensor must have contiguous last dimension");
-
-    CHECK_DEVICE(block_table);
-    TORCH_CHECK(block_table.dtype() == torch::kInt32, "block_table must have dtype torch.int32");
+    CHECK_CONTIGUOUS(seqlens_k);
     TORCH_CHECK(block_table.stride(-1) == 1, "block_table must have contiguous last dimension");
+    CHECK_CONTIGUOUS(tile_scheduler_metadata);
+    CHECK_CONTIGUOUS(num_splits);
 
     const auto sizes = q.sizes();
     const int batch_size = sizes[0];
     const int seqlen_q_ori = sizes[1];
-    const int num_heads_ori = sizes[2];
-    const int head_size = sizes[3];
-    TORCH_CHECK(head_size % 8 == 0, "head_size should be a multiple of 8");
-    TORCH_CHECK(head_size_v % 32 == 0, "head_size_v should be a multiple of 32");
+    const int num_heads_q = sizes[2];
+    const int head_size_k = sizes[3];
+    TORCH_CHECK(head_size_k == 576, "Only head_size_k == 576 is supported");
+    TORCH_CHECK(head_size_v == 512, "Only head_size_v == 576 is supported");
 
     const int max_num_blocks_per_seq = block_table.size(1);
     const int num_blocks = kcache.size(0);
     const int page_block_size = kcache.size(1);
     const int num_heads_k = kcache.size(2);
     TORCH_CHECK(batch_size > 0, "batch size must be postive");
-    TORCH_CHECK(num_heads_ori % num_heads_k == 0, "Number of heads in key/value must divide number of heads in query");
+    TORCH_CHECK(num_heads_q % num_heads_k == 0, "Number of heads in key/value must divide number of heads in query");
 
     if (seqlen_q_ori == 1) { is_causal = false; }
 
-    const int ngroups = num_heads_ori / num_heads_k;
-    const int seqlen_q = seqlen_q_ori * ngroups;
+    const int num_q_heads_per_hk = num_heads_q / num_heads_k;
+    const int q_seq_per_hk = seqlen_q_ori * num_q_heads_per_hk;
     const int num_heads = num_heads_k;
-    q = q.view({batch_size, seqlen_q_ori, num_heads_k, ngroups, head_size}).transpose(2, 3)
-            .reshape({batch_size, seqlen_q, num_heads, head_size});
+    q = q.view({batch_size, seqlen_q_ori, num_heads_k, num_q_heads_per_hk, head_size_k}).transpose(2, 3)
+            .reshape({batch_size, q_seq_per_hk, num_heads, head_size_k});
 
-    int head_size_k = head_size;
-    CHECK_SHAPE(q, batch_size, seqlen_q, num_heads, head_size);
+    CHECK_SHAPE(q, batch_size, q_seq_per_hk, num_heads, head_size_k);
     CHECK_SHAPE(kcache, num_blocks, page_block_size, num_heads_k, head_size_k);
-    if (vcache_.has_value()) { CHECK_SHAPE(vcache, num_blocks, page_block_size, num_heads_k, head_size_v); }
-    CHECK_SHAPE(block_table, batch_size, max_num_blocks_per_seq);
-
-
-    TORCH_CHECK(seqlens_k.dtype() == torch::kInt32, "seqlens_k must have dtype int32");
-    CHECK_DEVICE(seqlens_k);
-    CHECK_CONTIGUOUS(seqlens_k);
     CHECK_SHAPE(seqlens_k, batch_size);
+    CHECK_SHAPE(block_table, batch_size, max_num_blocks_per_seq);
+    TORCH_CHECK(tile_scheduler_metadata.size(1) == TileSchedulerMetaDataSize);
+    CHECK_SHAPE(num_splits, batch_size+1);
 
     at::cuda::CUDAGuard device_guard{(char)q.get_device()};
 
     auto opts = q.options();
-    at::Tensor out = torch::empty({batch_size, seqlen_q, num_heads, head_size_v}, opts);
-    at::Tensor softmax_lse = torch::empty({batch_size, num_heads, seqlen_q}, opts.dtype(at::kFloat));
+    at::Tensor out = torch::empty({batch_size, q_seq_per_hk, num_heads, head_size_v}, opts);
+    at::Tensor softmax_lse = torch::empty({batch_size, num_heads, q_seq_per_hk}, opts.dtype(at::kFloat));
+    CHECK_CONTIGUOUS(softmax_lse);
 
     Flash_fwd_mla_params params = {};
     // Set the sizes.
     params.b = batch_size;
-    params.seqlen_q = seqlen_q;
-    params.cu_seqlens_k = seqlens_k.data_ptr<int>();
-    params.h = num_heads;
-    params.h_h_k_ratio = num_heads / num_heads_k;
-    params.ngroups = ngroups;
+    params.s_q = seqlen_q_ori;
+    params.q_seq_per_hk = q_seq_per_hk;
+    params.seqlens_k_ptr = seqlens_k.data_ptr<int>();
+    params.h_q = num_heads_q;
+    params.h_k = num_heads_k;
+    params.num_blocks = num_blocks;
+    params.q_head_per_hk = num_q_heads_per_hk;
     params.is_causal = is_causal;
-    params.d = head_size;
+    params.d = head_size_k;
     params.d_v = head_size_v;
     params.scale_softmax = softmax_scale;
     params.scale_softmax_log2 = float(softmax_scale * M_LOG2E);
     // Set the pointers and strides.
     params.q_ptr = q.data_ptr();
     params.k_ptr = kcache.data_ptr();
-    params.v_ptr = vcache.data_ptr();
     params.o_ptr = out.data_ptr();
     params.softmax_lse_ptr = softmax_lse.data_ptr();
     // All stride are in elements, not bytes.
     params.q_batch_stride = q.stride(0);
     params.k_batch_stride = kcache.stride(0);
-    params.v_batch_stride = vcache.stride(0);
     params.o_batch_stride = out.stride(0);
     params.q_row_stride = q.stride(-3);
     params.k_row_stride = kcache.stride(-3);
-    params.v_row_stride = vcache.stride(-3);
     params.o_row_stride = out.stride(-3);
     params.q_head_stride = q.stride(-2);
     params.k_head_stride = kcache.stride(-2);
-    params.v_head_stride = vcache.stride(-2);
     params.o_head_stride = out.stride(-2);
 
     params.block_table = block_table.data_ptr<int>();
     params.block_table_batch_stride = block_table.stride(0);
     params.page_block_size = page_block_size;
-
-    TORCH_CHECK(tile_scheduler_metadata.dtype() == torch::kInt32, "tile_scheduler_metadata must have dtype int32");
-    TORCH_CHECK(tile_scheduler_metadata.size(1) == TileSchedulerMetaDataSize);
-    CHECK_DEVICE(tile_scheduler_metadata);
-    CHECK_CONTIGUOUS(tile_scheduler_metadata);
+    
     params.tile_scheduler_metadata_ptr = tile_scheduler_metadata.data_ptr<int>();
     params.num_sm_parts = tile_scheduler_metadata.size(0);
-    TORCH_CHECK(num_splits.dtype() == torch::kInt32, "num_splits must have dtype int32");
-    CHECK_DEVICE(num_splits);
-    CHECK_CONTIGUOUS(num_splits);
     params.num_splits_ptr = num_splits.data_ptr<int>();
 
-    at::Tensor softmax_lse_accum = torch::empty({batch_size + params.num_sm_parts, num_heads, seqlen_q}, opts.dtype(at::kFloat));
-    at::Tensor out_accum = torch::empty({batch_size + params.num_sm_parts, num_heads, seqlen_q, head_size_v}, opts.dtype(at::kFloat));
+    const int total_num_splits = batch_size + params.num_sm_parts;
+    at::Tensor softmax_lse_accum = torch::empty({total_num_splits, num_heads, q_seq_per_hk}, opts.dtype(at::kFloat));
+    at::Tensor out_accum = torch::empty({total_num_splits, num_heads, q_seq_per_hk, head_size_v}, opts.dtype(at::kFloat));
+    CHECK_CONTIGUOUS(softmax_lse_accum);
+    CHECK_CONTIGUOUS(out_accum);
+    params.total_num_splits = total_num_splits;
     params.softmax_lseaccum_ptr = softmax_lse_accum.data_ptr();
     params.oaccum_ptr = out_accum.data_ptr();
 
     auto stream = at::cuda::getCurrentCUDAStream().stream();
-    TORCH_CHECK(head_size == 576);
-
+    TORCH_CHECK(head_size_k == 576);
     if (q_dtype == torch::kBFloat16) {
-        run_mha_fwd_splitkv_mla<cutlass::bfloat16_t, 576>(params, stream);
-    }
-    #ifndef FLASH_MLA_DISABLE_FP16
-    else if (q_dtype == torch::kHalf) {
-        run_mha_fwd_splitkv_mla<cutlass::half_t, 576>(params, stream);
-    }
-    #endif
-    else {
+        run_flash_splitkv_mla_kernel<cutlass::bfloat16_t>(params, stream);
+        run_flash_mla_combine_kernel<cutlass::bfloat16_t>(params, stream);
+    } else if (q_dtype == torch::kHalf) {
+#ifdef FLASH_MLA_DISABLE_FP16
+        TORCH_CHECK(false, "FlashMLA is compiled with -DFLASH_MLA_DISABLE_FP16. Please remove this flag from your environment and re-compile FlashMLA.");
+#else
+        run_flash_splitkv_mla_kernel<cutlass::half_t>(params, stream);
+        run_flash_mla_combine_kernel<cutlass::half_t>(params, stream);
+#endif
+    } else {
         TORCH_CHECK(false, "Unsupported tensor dtype for query");
     }
 
-    out = out.view({batch_size, seqlen_q_ori, ngroups, num_heads_k, head_size_v}).transpose(2, 3)
-            .reshape({batch_size, seqlen_q_ori, num_heads_ori, head_size_v});
-    softmax_lse = softmax_lse.view({batch_size, num_heads_k, seqlen_q_ori, ngroups}).transpose(2, 3)
-            .reshape({batch_size, num_heads_ori, seqlen_q_ori});
+    out = out.view({batch_size, seqlen_q_ori, num_q_heads_per_hk, num_heads_k, head_size_v}).transpose(2, 3)
+            .reshape({batch_size, seqlen_q_ori, num_heads_q, head_size_v});
+    softmax_lse = softmax_lse.view({batch_size, num_heads_k, seqlen_q_ori, num_q_heads_per_hk}).transpose(2, 3)
+            .reshape({batch_size, num_heads_q, seqlen_q_ori});
 
     return {out, softmax_lse};
 }

csrc/flash_fwd_mla_bf16_sm90.cu

-#include "flash_fwd_mla_kernel.h"
-
-template void run_mha_fwd_splitkv_mla<cutlass::bfloat16_t, 576>(Flash_fwd_mla_params &params, cudaStream_t stream);

csrc/flash_fwd_mla_fp16_sm90.cu

-#include "flash_fwd_mla_kernel.h"
-
-template void run_mha_fwd_splitkv_mla<cutlass::half_t, 576>(Flash_fwd_mla_params &params, cudaStream_t stream);

csrc/kernels/config.h

+#pragma once
+
+namespace Config {
+
+static constexpr int BLOCK_SIZE_M = 64;
+static constexpr int PAGE_BLOCK_SIZE = 64;
+
+static constexpr int HEAD_DIM_K = 576;
+static constexpr int HEAD_DIM_V = 512;
+
+static constexpr int FIXED_OVERHEAD_NUM_BLOCKS = 5;
+
+}

csrc/flash_fwd_mla_metadata.cu → csrc/kernels/get_mla_metadata.cu

-#include "flash_fwd_mla_kernel.h"
+#include "get_mla_metadata.h"
 
-static constexpr int MaxBatchSize = 4096;
+#include <cuda_runtime_api.h>
+#include <cutlass/fast_math.h>
 
-__global__ void __launch_bounds__(256, 1, 1)
+#include "utils.h"
+
+__global__ void __launch_bounds__(32, 1, 1)
 get_mla_metadata_kernel(__grid_constant__ const Mla_metadata_params params) {
     int *seqlens_k_ptr = params.seqlens_k_ptr;
     int *tile_scheduler_metadata_ptr = params.tile_scheduler_metadata_ptr;
@@ -12,8 +15,9 @@ get_mla_metadata_kernel(__grid_constant__ const Mla_metadata_params params) {
     int fixed_overhead_num_blocks = params.fixed_overhead_num_blocks;
     int num_sm_parts = params.num_sm_parts;
 
-    __shared__ int num_blocks_shared[MaxBatchSize];
-    __shared__ int num_splits_shared[MaxBatchSize];
+    extern __shared__ int shared_mem[];
+    int* num_blocks_shared = shared_mem; // [batch_size]
+    int* num_splits_shared = shared_mem + batch_size; // [batch_size+1]
 
     int total_num_blocks = 0;
     for (int i = threadIdx.x; i < batch_size; i += 32) {
@@ -27,7 +31,7 @@ get_mla_metadata_kernel(__grid_constant__ const Mla_metadata_params params) {
     __syncwarp();
 
     if (threadIdx.x == 0) {
-        int payload = cutlass::ceil_div(total_num_blocks, num_sm_parts) + fixed_overhead_num_blocks;
+        int payload = max(cutlass::ceil_div(total_num_blocks, num_sm_parts) + fixed_overhead_num_blocks, 2*fixed_overhead_num_blocks);
 
         int now_idx = 0, now_block = 0, now_n_split_idx = 0, cum_num_splits = 0;
         num_splits_shared[0] = 0;
@@ -70,8 +74,9 @@ get_mla_metadata_kernel(__grid_constant__ const Mla_metadata_params params) {
     }
 }
 
-void get_mla_metadata_func(Mla_metadata_params &params, cudaStream_t stream) {
-    FLASH_ASSERT(params.batch_size < MaxBatchSize);
-    get_mla_metadata_kernel<<<1, 32, 0, stream>>>(params);
+void run_get_mla_metadata_kernel(Mla_metadata_params &params, cudaStream_t stream) {
+    int smem_size = sizeof(int) * (params.batch_size*2+1);
+    CHECK_CUDA(cudaFuncSetAttribute(get_mla_metadata_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    get_mla_metadata_kernel<<<1, 32, smem_size, stream>>>(params);
     CHECK_CUDA_KERNEL_LAUNCH();
-}
\ No newline at end of file
+}

csrc/kernels/get_mla_metadata.h

+#pragma once
+
+#include "params.h"
+
+void run_get_mla_metadata_kernel(Mla_metadata_params &params, cudaStream_t stream);

csrc/kernels/mla_combine.cu

+#include "mla_combine.h"
+
+#include <cute/tensor.hpp>
+#include <cutlass/cutlass.h>
+#include <cutlass/array.h>
+#include <cutlass/numeric_types.h>
+
+#include "params.h"
+#include "utils.h"
+#include "config.h"  // for BLOCK_SIZE_M and HEAD_DIM_V
+
+using namespace cute;
+
+template<typename ElementT, int HEAD_DIM_V, int BLOCK_SIZE_M, int MAX_SPLITS, int NUM_THREADS>
+__global__ void __launch_bounds__(NUM_THREADS)
+flash_fwd_mla_combine_kernel(__grid_constant__ const Flash_fwd_mla_params params) {
+    // grid_shape: [batch_size, num_q_heads*s_q / BLOCK_SIZE_M]
+    // Each CTA gathers the activation of some heads from one batch, do scaling & accumulation, and save the result
+    static_assert(NUM_THREADS/32 == BLOCK_SIZE_M); // The number of warps == block_size_m
+    const int batch_idx = blockIdx.x;
+    const int m_block_idx = blockIdx.y;
+    const int warp_idx = threadIdx.x / 32;
+    const int lane_idx = threadIdx.x % 32;
+
+    const int start_split_idx = __ldg(params.num_splits_ptr + batch_idx);
+    const int end_split_idx = __ldg(params.num_splits_ptr + batch_idx + 1);
+    const int my_num_splits = end_split_idx - start_split_idx;
+    FLASH_DEVICE_ASSERT(my_num_splits <= MAX_SPLITS);
+    if (my_num_splits == 1) {
+        return;
+    }
+    
+    const int num_q_seqs = params.q_seq_per_hk * params.h_k;
+    const int num_cur_valid_q_seqs = min(BLOCK_SIZE_M, num_q_seqs - m_block_idx*BLOCK_SIZE_M);
+    Tensor gLseAccum = make_tensor(
+        make_gmem_ptr((float*)params.softmax_lseaccum_ptr + start_split_idx*num_q_seqs + m_block_idx*BLOCK_SIZE_M),
+        Shape<Int<MAX_SPLITS>, Int<BLOCK_SIZE_M>>{},
+        make_stride(num_q_seqs, _1{})
+    );
+    Tensor gLse = make_tensor(
+        make_gmem_ptr((float*)params.softmax_lse_ptr + batch_idx*num_q_seqs + m_block_idx*BLOCK_SIZE_M),
+        Shape<Int<BLOCK_SIZE_M>>{},
+        Stride<_1>{}
+    );
+    
+    extern __shared__ float smem_buf[];
+    Tensor sLseScale = make_tensor(
+        make_smem_ptr(smem_buf),
+        Shape<Int<BLOCK_SIZE_M>, Int<MAX_SPLITS>>{},
+        Stride<Int<MAX_SPLITS+1>, _1>{} // +1 to avoid bank conflict
+    );
+    
+    // Wait for the previous kernel (the MLA kernel) to finish
+    cudaGridDependencySynchronize();
+    
+    // Read gLseAccum into sLseScale
+    {
+        #pragma unroll 4
+        for (int elem_idx = threadIdx.x; elem_idx < my_num_splits*BLOCK_SIZE_M; elem_idx += NUM_THREADS) {
+            int split_idx = elem_idx / BLOCK_SIZE_M;
+            int seq_idx = elem_idx % BLOCK_SIZE_M;
+            sLseScale(seq_idx, split_idx) = seq_idx < num_cur_valid_q_seqs ? gLseAccum(split_idx, seq_idx) : -INFINITY;
+        }
+        __syncthreads();
+    }
+
+    if (warp_idx >= num_cur_valid_q_seqs)
+        return;
+
+    // Warp #i gathers LseAccum for seq #i
+    {
+        constexpr int NUM_LSE_PER_THREAD = cute::ceil_div(MAX_SPLITS, 32);
+        float local_lse[NUM_LSE_PER_THREAD];
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < NUM_LSE_PER_THREAD; ++i) {
+            const int split_idx = i*32 + lane_idx;
+            local_lse[i] = split_idx < my_num_splits ? sLseScale(warp_idx, split_idx) : -INFINITY;
+        }
+
+        float max_lse = -INFINITY;
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < NUM_LSE_PER_THREAD; ++i)
+            max_lse = max(max_lse, local_lse[i]);
+        CUTLASS_PRAGMA_UNROLL
+        for (int offset = 16; offset >= 1; offset /= 2)
+            max_lse = max(max_lse, __shfl_xor_sync(uint32_t(-1), max_lse, offset));
+        max_lse = max_lse == -INFINITY ? 0.0f : max_lse;  // In case all local LSEs are -inf
+
+        float sum_lse = 0;
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < NUM_LSE_PER_THREAD; ++i)
+            sum_lse = sum_lse + exp2f(local_lse[i] - max_lse);
+        CUTLASS_PRAGMA_UNROLL
+        for (int offset = 16; offset >= 1; offset /= 2)
+            sum_lse = sum_lse + __shfl_xor_sync(uint32_t(-1), sum_lse, offset);
+
+        float global_lse = (sum_lse == 0.f || sum_lse != sum_lse) ? INFINITY : log2f(sum_lse) + max_lse;
+        if (lane_idx == 0)
+            gLse(warp_idx) = global_lse / (float)M_LOG2E;
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < NUM_LSE_PER_THREAD; ++i) {
+            const int split_idx = i*32 + lane_idx;
+            if (split_idx < my_num_splits) sLseScale(warp_idx, split_idx) = exp2f(local_lse[i] - global_lse);
+        }
+    }
+
+    __syncwarp();
+
+    // Warp #i accumulates activation for seq #i
+    {
+        const int64_t row_offset_oaccum = (int64_t)(start_split_idx*num_q_seqs+m_block_idx*BLOCK_SIZE_M+warp_idx) * HEAD_DIM_V;
+        Tensor gOaccum = make_tensor(
+            make_gmem_ptr(reinterpret_cast<float *>(params.oaccum_ptr) + row_offset_oaccum),
+            Shape<Int<MAX_SPLITS>, Int<HEAD_DIM_V>>{},
+            make_stride(num_q_seqs*HEAD_DIM_V, _1{})
+        );
+
+        static_assert(HEAD_DIM_V % 32 == 0);
+        constexpr int ELEMS_PER_THREAD = HEAD_DIM_V / 32;
+        float result[ELEMS_PER_THREAD];
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < ELEMS_PER_THREAD; ++i)
+            result[i] = 0.0f;
+
+        #pragma unroll 2
+        for (int split = 0; split < my_num_splits; ++split) {
+            float lse_scale = sLseScale(warp_idx, split);
+            if (lse_scale != 0.f) {
+                CUTLASS_PRAGMA_UNROLL
+                for (int i = 0; i < ELEMS_PER_THREAD; ++i) {
+                    result[i] += lse_scale * gOaccum(split, lane_idx + i*32);
+                }
+            }
+        }
+
+        cudaTriggerProgrammaticLaunchCompletion();
+        
+        const int q_seq_idx = m_block_idx*BLOCK_SIZE_M + warp_idx;
+        const int k_head_idx = q_seq_idx / params.q_seq_per_hk;
+        auto o_ptr = reinterpret_cast<ElementT *>(params.o_ptr) + batch_idx*params.o_batch_stride + k_head_idx*params.o_head_stride + (q_seq_idx%params.q_seq_per_hk)*params.o_row_stride;
+        Tensor gO = make_tensor(
+            make_gmem_ptr(o_ptr),
+            Shape<Int<HEAD_DIM_V>>{},
+            Stride<_1>{}
+        );
+
+        CUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < ELEMS_PER_THREAD; ++i)
+            gO(lane_idx+i*32) = (ElementT)result[i];
+    }
+}
+
+
+#define MLA_NUM_SPLITS_SWITCH(NUM_SPLITS, NAME, ...)       \
+    [&] {                                                  \
+        if (NUM_SPLITS <= 32) {                            \
+            constexpr static int NAME = 32;                \
+            return __VA_ARGS__();                          \
+        } else if (NUM_SPLITS <= 64) {                     \
+            constexpr static int NAME = 64;                \
+            return __VA_ARGS__();                          \
+        } else if (NUM_SPLITS <= 96) {                     \
+            constexpr static int NAME = 96;                \
+            return __VA_ARGS__();                          \
+        } else if (NUM_SPLITS <= 128) {                    \
+            constexpr static int NAME = 128;               \
+            return __VA_ARGS__();                          \
+        } else if (NUM_SPLITS <= 160) {                    \
+            constexpr static int NAME = 160;               \
+            return __VA_ARGS__();                          \
+        } else {                                           \
+            FLASH_ASSERT(false);                           \
+        }                                                  \
+    }()
+
+
+template<typename ElementT>
+void run_flash_mla_combine_kernel(Flash_fwd_mla_params &params, cudaStream_t stream) {
+    MLA_NUM_SPLITS_SWITCH(params.num_sm_parts, NUM_SPLITS, [&] {
+        constexpr int BLOCK_SIZE_M = 8;
+        constexpr int NUM_THREADS = BLOCK_SIZE_M*32;
+        constexpr size_t smem_size = BLOCK_SIZE_M*(NUM_SPLITS+1)*sizeof(float);
+        auto combine_kernel = &flash_fwd_mla_combine_kernel<ElementT, Config::HEAD_DIM_V, BLOCK_SIZE_M, NUM_SPLITS, NUM_THREADS>;
+        CHECK_CUDA(cudaFuncSetAttribute(combine_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+        // Use cudaLaunchKernelEx to enable PDL (Programmatic Dependent Launch)
+        cudaLaunchAttribute attribute[1];
+        attribute[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+        attribute[0].val.programmaticStreamSerializationAllowed = 1;
+        cudaLaunchConfig_t combine_kernel_config = {
+            dim3(params.b, cute::ceil_div(params.h_k*params.q_seq_per_hk, BLOCK_SIZE_M), 1),
+            dim3(NUM_THREADS, 1, 1),
+            smem_size,
+            stream,
+            attribute,
+            1
+        };
+        cudaLaunchKernelEx(&combine_kernel_config, combine_kernel, params);
+    });
+    CHECK_CUDA_KERNEL_LAUNCH();
+}
+
+template void run_flash_mla_combine_kernel<cutlass::bfloat16_t>(Flash_fwd_mla_params &params, cudaStream_t stream);
+
+#ifndef FLASH_MLA_DISABLE_FP16
+template void run_flash_mla_combine_kernel<cutlass::half_t>(Flash_fwd_mla_params &params, cudaStream_t stream);
+#endif
\ No newline at end of file

csrc/kernels/mla_combine.h

+#pragma once
+
+#include "params.h"
+
+template<typename ElementT>
+void run_flash_mla_combine_kernel(Flash_fwd_mla_params &params, cudaStream_t stream);

csrc/flash_mla.h → csrc/kernels/params.h

@@ -5,39 +5,41 @@
 struct Flash_fwd_mla_params {
     using index_t = int64_t;
 
-    int b, seqlen_q, d, d_v;
-    int h, h_h_k_ratio, ngroups;
+    int b;              // batch size
+    int s_q;
+    int q_seq_per_hk;   // The number of q(s) per KV head, = h_q / h_k * s_q
+    int d, d_v;         // K/V dimension
+    int h_q, h_k;       // The number of Q/K heads
+    int num_blocks;     // Number of blocks in total
+    int q_head_per_hk;  // The number of q_head(s) per KV head, = h_q / h_k
     bool is_causal;
     float scale_softmax, scale_softmax_log2;
-    int *__restrict__ cu_seqlens_k;
-
+    
     void *__restrict__ q_ptr;
     void *__restrict__ k_ptr;
-    void *__restrict__ v_ptr;
     void *__restrict__ o_ptr;
     void *__restrict__ softmax_lse_ptr;
 
     index_t q_batch_stride;
     index_t k_batch_stride;
-    index_t v_batch_stride;
     index_t o_batch_stride;
     index_t q_row_stride;
     index_t k_row_stride;
-    index_t v_row_stride;
     index_t o_row_stride;
     index_t q_head_stride;
     index_t k_head_stride;
-    index_t v_head_stride;
     index_t o_head_stride;
 
     int *__restrict__ block_table;
     index_t block_table_batch_stride;
     int page_block_size;
+    int *__restrict__ seqlens_k_ptr;
 
     int *__restrict__ tile_scheduler_metadata_ptr;
     int num_sm_parts;
     int *__restrict__ num_splits_ptr;
 
+    int total_num_splits;
     void *__restrict__ softmax_lseaccum_ptr;
     void *__restrict__ oaccum_ptr;
 };
@@ -45,11 +47,6 @@ struct Flash_fwd_mla_params {
 static constexpr int TileSchedulerMetaDataSize = 8;
 // [begin_idx, begin_seqlen, end_idx, end_seqlen, begin_n_split_idx, _, _, _]
 
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template<typename T, int Headdim>
-void run_mha_fwd_splitkv_mla(Flash_fwd_mla_params &params, cudaStream_t stream);
-
 struct Mla_metadata_params {
     int *__restrict__ seqlens_k_ptr;
     int *__restrict__ tile_scheduler_metadata_ptr;
@@ -59,5 +56,3 @@ struct Mla_metadata_params {
     int fixed_overhead_num_blocks;
     int num_sm_parts;
 };
-
-void get_mla_metadata_func(Mla_metadata_params &params, cudaStream_t stream);

csrc/kernels/splitkv_mla.h

+#pragma once
+
+#include "params.h"
+
+template<typename InputT>
+void run_flash_splitkv_mla_kernel(Flash_fwd_mla_params &params, cudaStream_t stream);

csrc/kernels/traits.h

+#pragma once
+
+#include <cute/tensor.hpp>
+#include <cutlass/cutlass.h>
+#include <cutlass/numeric_types.h>
+#include <cutlass/barrier.h>
+
+#include "config.h"
+
+using TMABarrier = cutlass::arch::ClusterTransactionBarrier;
+using namespace cute;
+
+template<typename InputT_>
+struct Traits {
+    using InputT = InputT_;
+    
+    static constexpr int BLOCK_SIZE_M = Config::BLOCK_SIZE_M;
+    static constexpr int PAGE_BLOCK_SIZE = Config::PAGE_BLOCK_SIZE;
+    static constexpr int HEAD_DIM_K = Config::HEAD_DIM_K;
+    static constexpr int HEAD_DIM_V = Config::HEAD_DIM_V;
+
+    static constexpr int NUM_THREADS = 256;
+
+    static_assert(std::is_same_v<InputT, cutlass::bfloat16_t> || std::is_same_v<InputT, cutlass::half_t>);
+
+    using TiledMMA_QK_sQ = decltype(make_tiled_mma(
+        GMMA::ss_op_selector<InputT, InputT, float, Shape<Int<BLOCK_SIZE_M>, Int<PAGE_BLOCK_SIZE>, Int<HEAD_DIM_K>>, GMMA::Major::K, GMMA::Major::K>(),
+        Layout<Shape<_1, _1, _1>>{}
+    ));
+
+    using TiledMMA_QK_rQ = decltype(make_tiled_mma(
+        GMMA::rs_op_selector<InputT, InputT, float, Shape<Int<BLOCK_SIZE_M>, Int<PAGE_BLOCK_SIZE>, Int<HEAD_DIM_K>>, GMMA::Major::K, GMMA::Major::K>(),
+        Layout<Shape<_1, _1, _1>>{}
+    ));
+
+    using TiledMMA_PV_LocalP = decltype(make_tiled_mma(
+        GMMA::rs_op_selector<InputT, InputT, float, Shape<Int<BLOCK_SIZE_M>, Int<HEAD_DIM_V/2>, Int<PAGE_BLOCK_SIZE>>, GMMA::Major::K, GMMA::Major::MN>(),
+        Layout<Shape<_1, _1, _1>>{}
+    ));
+
+    using TiledMMA_PV_RemoteP = decltype(make_tiled_mma(
+        GMMA::ss_op_selector<InputT, InputT, float, Shape<Int<BLOCK_SIZE_M>, Int<HEAD_DIM_V/2>, Int<PAGE_BLOCK_SIZE>>, GMMA::Major::K, GMMA::Major::MN>(),
+        Layout<Shape<_1, _1, _1>>{}
+    ));
+
+    using SmemLayoutQ = decltype(tile_to_shape(
+        GMMA::Layout_K_SW128_Atom<InputT>{},
+        Shape<Int<BLOCK_SIZE_M>, Int<HEAD_DIM_K>>{}
+    ));
+
+    using SmemLayoutK = decltype(tile_to_shape(
+        GMMA::Layout_K_SW128_Atom<InputT>{},
+        Shape<Int<PAGE_BLOCK_SIZE>, Int<HEAD_DIM_K>>{}
+    ));
+
+    using SmemLayoutV = decltype(composition(
+        SmemLayoutK{},
+        make_layout(Shape<Int<HEAD_DIM_V>, Int<PAGE_BLOCK_SIZE>>{}, GenRowMajor{})
+    ));	// A transposed version of SmemLayoutK
+
+    using SmemLayoutP0 = decltype(tile_to_shape(
+        GMMA::Layout_K_SW128_Atom<InputT>{},
+        Shape<Int<BLOCK_SIZE_M>, Int<PAGE_BLOCK_SIZE>>{}
+    ));
+
+    using rP0Layout = decltype(layout(partition_fragment_C(
+        TiledMMA_QK_sQ{},
+        Shape<Int<BLOCK_SIZE_M>, Int<PAGE_BLOCK_SIZE>>{}
+    )));
+
+    struct SharedMemoryPlan {
+        cute::array_aligned<InputT, cosize_v<SmemLayoutQ>> smem_sQ;
+        cute::array_aligned<InputT, cosize_v<SmemLayoutK>> smem_sK0;
+        cute::array_aligned<InputT, cosize_v<SmemLayoutK>> smem_sK1;
+        cute::array_aligned<InputT, cosize_v<SmemLayoutP0>> smem_sP0;
+        cute::array_aligned<float, BLOCK_SIZE_M> smem_sM;
+        cute::array_aligned<float, 2*BLOCK_SIZE_M> sL_reduction_wksp;
+        cute::array_aligned<float, BLOCK_SIZE_M> smem_sScale0;
+        cute::array_aligned<float, BLOCK_SIZE_M> smem_sScale1;
+        TMABarrier barriers_K0[HEAD_DIM_K/64];
+        TMABarrier barriers_K1[HEAD_DIM_K/64];
+        TMABarrier barrier_Q;
+    };
+
+};
+
+template<
+    typename ShapeQ, typename TMA_Q,
+    typename ShapeK, typename TMA_K,
+    typename ShapeO, typename TMA_O
+>
+struct TmaParams {
+    ShapeQ shape_Q;
+    TMA_Q tma_Q;
+    ShapeK shape_K;
+    TMA_K tma_K;
+    ShapeO shape_O;
+    TMA_O tma_O;
+};
+
+enum NamedBarriers : int {
+    sScale0Ready = 0,
+    sScale1Ready = 1,
+    sP0Ready = 2,
+    rO1sP0sV0RIssued = 3
+};

csrc/static_switch.h → csrc/kernels/utils.h

@@ -5,7 +5,7 @@
         cudaError_t status_ = call;                                                                       \
         if (status_ != cudaSuccess) {                                                                     \
             fprintf(stderr, "CUDA error (%s:%d): %s\n", __FILE__, __LINE__, cudaGetErrorString(status_)); \
-            exit(1);                                                                                      \
+            exit(1);                                                                              \
         }                                                                                                 \
     } while(0)
 
@@ -29,37 +29,4 @@
         }                                                                                                 \
     } while(0)
 
-
-#define BOOL_SWITCH(COND, CONST_NAME, ...)      \
-  [&] {                                         \
-    if (COND) {                                 \
-      constexpr static bool CONST_NAME = true;  \
-      return __VA_ARGS__();                     \
-    } else {                                    \
-      constexpr static bool CONST_NAME = false; \
-      return __VA_ARGS__();                     \
-    }                                           \
-  }()
-
-
-#define MLA_NUM_SPLITS_SWITCH(NUM_SPLITS, NAME, ...) \
-  [&] {                                              \
-    if (NUM_SPLITS <= 32) {                          \
-      constexpr static int NAME = 32;                \
-      return __VA_ARGS__();                          \
-    } else if (NUM_SPLITS <= 64) {                   \
-      constexpr static int NAME = 64;                \
-      return __VA_ARGS__();                          \
-    } else if (NUM_SPLITS <= 96) {                   \
-      constexpr static int NAME = 96;                \
-      return __VA_ARGS__();                          \
-    } else if (NUM_SPLITS <= 128) {                  \
-      constexpr static int NAME = 128;               \
-      return __VA_ARGS__();                          \
-    } else if (NUM_SPLITS <= 160) {                  \
-      constexpr static int NAME = 160;               \
-      return __VA_ARGS__();                          \
-    } else {                                         \
-      FLASH_ASSERT(false);                           \
-    }                                                \
-  }()
+#define println(fmt, ...) { print(fmt, ##__VA_ARGS__); print("\n"); }

csrc/named_barrier.h

-#pragma once
-
-#include "cutlass/barrier.h"
-
-namespace flash {
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-// Enumerates the reserved named barriers to avoid potential conflicts
-
-enum class NamedBarriers {
-    SReady = 1,
-    SoftmaxReady = 2,
-};
-
-} // flash

csrc/softmax.h

-// Adapted from https://github.com/Dao-AILab/flash-attention/blob/main/csrc/flash_attn/src/softmax.h
-/******************************************************************************
- * Copyright (c) 2024, Tri Dao.
- ******************************************************************************/
-
-#pragma once
-
-#include <cmath>
-
-#include <cute/tensor.hpp>
-#include <cutlass/numeric_types.h>
-
-#include "utils.h"
-
-namespace flash {
-
-using namespace cute;
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template<bool zero_init=true, 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) = zero_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 zero_init=true, 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_<zero_init>(tensor, summary, op);
-    quad_allreduce_(summary, summary, op);
-}
-
-template<bool zero_init=true, 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_<zero_init>(tensor, max, max_op);
-}
-
-template<bool zero_init=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_<zero_init>(tensor, sum, sum_op);
-}
-
-// Apply the exp to all the elements.
-template <bool Scale_max=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__forceinline__ __device__ auto 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) {
-        // If max is -inf, then all elements must have been -inf (possibly due to masking).
-        // We don't want (-inf - (-inf)) since that would give NaN.
-        // If we don't have float around M_LOG2E the multiplication is done in fp64.
-        const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * (Scale_max ? scale : float(M_LOG2E));
-        #pragma unroll
-        for (int ni = 0; ni < size<1>(tensor); ++ni)  {
-            // Instead of computing exp(x - max), we compute exp2(x * log_2(e) -
-            // max * log_2(e)) This allows the compiler to use the ffma
-            // instruction instead of fadd and fmul separately.
-            // The following macro will disable the use of fma.
-            // See: https://github.com/pytorch/pytorch/issues/121558 for more details
-            // This macro is set in PyTorch and not FlashAttention
-            #ifdef UNFUSE_FMA
-                tensor(mi, ni) = exp2f(__fmul_rn(tensor(mi, ni), scale) - max_scaled);
-            #else
-                tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
-            #endif
-        }
-    }
-    return tensor;
-}
-
-// Apply the exp to all the elements.
-template <bool zero_init=true, typename Engine0, typename Layout0, typename Engine1, typename Layout1>
-__forceinline__ __device__ void max_scale_exp2_sum(Tensor<Engine0, Layout0> &tensor, Tensor<Engine1, Layout1> &max, Tensor<Engine1, Layout1> &sum, 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) {
-        MaxOp<float> max_op;
-        max(mi) = zero_init ? tensor(mi, 0) : max_op(max(mi), tensor(mi, 0));
-        #pragma unroll
-        for (int ni = 1; ni < size<1>(tensor); ni++) {
-            max(mi) = max_op(max(mi), tensor(mi, ni));
-        }
-        max(mi) = Allreduce<4>::run(max(mi), max_op);
-        // If max is -inf, then all elements must have been -inf (possibly due to masking).
-        // We don't want (-inf - (-inf)) since that would give NaN.
-        const float max_scaled = max(mi) == -INFINITY ? 0.f : max(mi) * scale;
-        sum(mi) = 0;
-        #pragma unroll
-        for (int ni = 0; ni < size<1>(tensor); ++ni)  {
-            // Instead of computing exp(x - max), we compute exp2(x * log_2(e) -
-            // max * log_2(e)) This allows the compiler to use the ffma
-            // instruction instead of fadd and fmul separately.
-            tensor(mi, ni) = exp2f(tensor(mi, ni) * scale - max_scaled);
-            sum(mi) += tensor(mi, ni);
-        }
-        SumOp<float> sum_op;
-        sum(mi) = Allreduce<4>::run(sum(mi), sum_op);
-    }
-}
-
-template<typename Tensor0, typename Tensor1>
-__forceinline__ __device__ void rescale_o(Tensor0 &acc_o, Tensor1 &scale_o) {
-    // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))
-    Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout()));
-    #pragma unroll
-    for (int mi = 0; mi < size(scale_o); ++mi) {
-        #pragma unroll
-        for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { acc_o_rowcol(mi, ni) *= scale_o(mi); }
-    }
-}
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template <int kNRows>
-struct Softmax {
-
-    using TensorT = decltype(make_tensor<float>(Shape<Int<kNRows>>{}));
-    TensorT row_max, row_sum;
-
-    __forceinline__ __device__ Softmax() {};
-
-    template<bool Is_first, bool Check_inf=false, typename Tensor0>
-    __forceinline__ __device__ TensorT softmax(Tensor0 &acc_s, float softmax_scale_log2) {
-        // 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(), flash::convert_layout_acc_rowcol(acc_s.layout()));
-        static_assert(decltype(size<0>(scores))::value == kNRows);
-        TensorT scale_o;
-        clear(scale_o);
-        if (Is_first) {
-            flash::template reduce_max</*zero_init=*/true>(scores, row_max);
-            flash::scale_apply_exp2(scores, row_max, softmax_scale_log2);
-            flash::reduce_sum</*zero_init=*/true>(scores, row_sum);
-        } else {
-            Tensor scores_max_prev = make_fragment_like(row_max);
-            cute::copy(row_max, scores_max_prev);
-            flash::template reduce_max</*zero_init=*/false>(scores, row_max);
-            // Reshape acc_o from (MMA=4, MMA_M, MMA_K) to (nrow=(2, MMA_M), ncol=(2, MMA_K))
-            #pragma unroll
-            for (int mi = 0; mi < size(row_max); ++mi) {
-                float scores_max_cur = !Check_inf
-                    ? row_max(mi)
-                    : (row_max(mi) == -INFINITY ? 0.0f : row_max(mi));
-                float scores_scale = exp2f((scores_max_prev(mi) - scores_max_cur) * softmax_scale_log2);
-                scale_o(mi) = scores_scale;
-                row_sum(mi) *= scores_scale;
-            }
-            flash::scale_apply_exp2(scores, row_max, softmax_scale_log2);
-            // We don't do the reduce across threads here since we don't need to use the row_sum.
-            // We do that reduce at the end when we need to normalize the softmax.
-            flash::reduce_sum</*zero_init=*/false>(scores, row_sum);
-        }
-        return scale_o;
-    };
-
-    template<bool Is_dropout=false, bool Split=false, typename Tensor0>
-    __forceinline__ __device__ TensorT normalize_softmax_lse(Tensor0 &acc_o, float softmax_scale, float rp_dropout=1.0) {
-        SumOp<float> sum_op;
-        quad_allreduce_(row_sum, row_sum, sum_op);
-        TensorT lse = make_fragment_like(row_sum);
-        // Reshape acc_s from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))
-        Tensor acc_o_rowcol = make_tensor(acc_o.data(), flash::convert_layout_acc_rowcol(acc_o.layout()));
-        static_assert(decltype(size<0>(acc_o_rowcol))::value == kNRows);
-        #pragma unroll
-        for (int mi = 0; mi < size<0>(acc_o_rowcol); ++mi) {
-            float sum = row_sum(mi);
-            float inv_sum = (sum == 0.f || sum != sum) ? 1.f : 1.f / sum;
-            lse(mi) = (sum == 0.f || sum != sum) ? (Split ? -INFINITY : INFINITY) : row_max(mi) * softmax_scale + __logf(sum);
-            float scale = !Is_dropout ? inv_sum : inv_sum * rp_dropout;
-            #pragma unroll
-            for (int ni = 0; ni < size<1>(acc_o_rowcol); ++ni) { acc_o_rowcol(mi, ni) *= scale; }
-        }
-        return lse;
-    };
-};
-
-}  // namespace flash

csrc/utils.h

-// Adapted from https://github.com/Dao-AILab/flash-attention/blob/main/hopper/utils.h
-/******************************************************************************
- * Copyright (c) 2024, Tri Dao.
- ******************************************************************************/
-
-#pragma once
-
-#include <assert.h>
-#include <stdint.h>
-#include <stdlib.h>
-
-#include <cuda_bf16.h>
-
-#include <cute/tensor.hpp>
-
-#include <cutlass/array.h>
-#include <cutlass/cutlass.h>
-#include <cutlass/numeric_conversion.h>
-#include <cutlass/numeric_types.h>
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-namespace flash {
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-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 zero_init=false, int wg_wait=0, bool arrive=true, bool commit=true, typename Tensor0, typename Tensor1, typename Tensor2, typename TiledMma>
-__forceinline__ __device__ void gemm(TiledMma &tiled_mma, Tensor0 const &tCrA, Tensor1 const &tCrB, Tensor2 &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) { cute::warpgroup_fence_operand(const_cast<Tensor0 &>(tCrA)); }
-    warpgroup_fence_operand(tCrC);
-    if constexpr (arrive) {
-        warpgroup_arrive();
-    }
-    if constexpr (zero_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;
-        }
-    }
-    if constexpr (commit) {
-        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<Tensor0 &>(tCrA)); }
-}
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-// For SM80, convert acc_layout from (MMA=4, MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, MMA_N))
-// For SM90, convert acc_layout from ((2, 2, V), MMA_M, MMA_N) to (nrow=(2, MMA_M), ncol=(2, V, MMA_N))
-template<bool Transposed=false, typename Layout0>
-__forceinline__ __device__ auto convert_layout_acc_rowcol(Layout0 acc_layout) {
-    if constexpr (decltype(rank<0>(acc_layout))::value == 3) {  // SM90
-        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;
-        if constexpr (!Transposed) {
-            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)));
-        } else {
-             return make_layout(make_layout(get<0, 0>(l), get<0, 2>(l), get<2>(l)), make_layout(get<0, 1>(l), get<1>(l)));
-        }
-
-    } else {  // SM80
-        static_assert(decltype(size<0>(acc_layout))::value == 4);
-        static_assert(decltype(rank(acc_layout))::value == 3);
-        auto l = logical_divide(acc_layout, Shape<_2>{});  // ((2, 2), MMA_M, MMA_N)
-        if constexpr (!Transposed) {
-            return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));
-        } else {
-            return make_layout(make_layout(get<0, 0>(l), get<2>(l)), make_layout(get<0, 1>(l), get<1>(l)));
-        }
-    }
-};
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-// For SM80, convert acc_layout from (MMA=4, MMA_M, MMA_N) to ((4, 2), MMA_M, MMA_N / 2)
-// if using m16n8k16, or to (4, MMA_M, MMA_N) if using m16n8k8.
-// For SM90, FP16/BF16, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((2, 2, 2), MMA_M, (N / 16, MMA_N))
-// For SM90, FP8, convert acc_layout from ((2, 2, N / 8), MMA_M, MMA_N) to ((4, 2, 2), MMA_M, (N / 32, MMA_N))
-template<typename MMA_Traits, typename Layout0>
-__forceinline__ __device__ auto convert_layout_acc_Aregs(Layout0 acc_layout) {
-    using X = Underscore;
-    if constexpr (decltype(rank<0>(acc_layout))::value == 3) {  // SM90
-        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);
-        if constexpr (sizeof(typename MMA_Traits::ValTypeA) == 2) {
-            auto l = logical_divide(get<0, 2>(acc_layout), Tile<_2>{});  // ((2, N / 16))
-            return make_layout(make_layout(get<0, 0>(acc_layout), get<0, 1>(acc_layout), get<0, 0>(l)), get<1>(acc_layout), coalesce(make_layout(get<0, 1>(l), get<2>(acc_layout))));
-        } else {
-            static_assert(sizeof(typename MMA_Traits::ValTypeA) == 1);
-            static_assert(decltype(stride<0, 0>(acc_layout))::value == 1);
-            static_assert(decltype(stride<0, 1>(acc_layout))::value == 2);
-            auto l = logical_divide(get<0, 2>(acc_layout), Tile<Layout<Shape<_2, _2>>>{});  // (((2, 2), N / 32))
-            // This combines the first two modes (<0, 0> and <0, 1>) into one mode.
-            // Will require register shuffling later to be correct.
-            return make_layout(make_layout(Layout<_4>{}, get<0, 0, 0>(l), get<0, 0, 1>(l)),
-                               get<1>(acc_layout),
-                               coalesce(make_layout(get<0, 1>(l), get<2>(acc_layout))));  // ((4, 2, 2), MMA_M, N / 32 * MMA_N)
-            // This combination is right but doesn't work with register shuffling.
-            // return make_layout(make_layout(coalesce(make_layout(get<0, 0>(acc_layout), get<0, 0, 0>(l))), get<0, 1>(acc_layout), get<0, 0, 1>(l)),
-            //                    get<1>(acc_layout),
-            //                    coalesce(make_layout(get<0, 1>(l), get<2>(acc_layout))));
-        }
-    } else {  // SM80
-        static_assert(decltype(size<0>(acc_layout))::value == 4);
-        static_assert(decltype(rank(acc_layout))::value == 3);
-        constexpr int mma_shape_K = get<2>(typename MMA_Traits::Shape_MNK{});
-        static_assert(mma_shape_K == 8 || mma_shape_K == 16);
-        if constexpr (mma_shape_K == 8) {
-            return acc_layout;
-        } else {
-            auto l = logical_divide(acc_layout, Shape<X, X, _2>{});  // (4, MMA_M, (2, MMA_N / 2)))
-            return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
-        }
-    }
-};
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-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;
-    // HACK: this requires tensor to be "contiguous"
-    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());
-}
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-// Blocks until all but N previous cp.async.commit_group operations have committed.
-// This differs from cute::cp_async_wait in that when N = 0 we don't call cp.async.wait_all
-// (which is equivalent to commit_group then wait_group 0).
-// Instead we just call cp.async.wait_group 0, which is slightly faster.
-// https://github.com/NVIDIA/cutlass/blob/master/include/cute/arch/copy_sm80.hpp#L113
-template <int N>
-CUTE_HOST_DEVICE
-void cp_async_wait() {
-#if defined(CUTE_ARCH_CP_ASYNC_SM80_ENABLED)
-    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
-#endif
-}
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-template <bool Is_even_MN=true, bool Is_even_K=true, bool Clear_OOB_MN=false, bool Clear_OOB_K=true,
-          typename TiledCopy, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
-          typename Engine2, typename Layout2, typename Engine3, typename Layout3>
-__forceinline__ __device__ void copy(TiledCopy tiled_copy, Tensor<Engine0, Layout0> const &S,
-                            Tensor<Engine1, Layout1> &D, Tensor<Engine2, Layout2> const &identity_MN,
-                            Tensor<Engine3, Layout3> const &predicate_K, const int max_MN=0) {
-    CUTE_STATIC_ASSERT_V(rank(S) == Int<3>{});
-    CUTE_STATIC_ASSERT_V(rank(D) == Int<3>{});
-    CUTE_STATIC_ASSERT_V(size<0>(S) == size<0>(D));                     // MMA
-    CUTE_STATIC_ASSERT_V(size<1>(S) == size<1>(D));                     // MMA_M
-    CUTE_STATIC_ASSERT_V(size<2>(S) == size<2>(D));                     // MMA_K
-    // There's no case where !Clear_OOB_K && Clear_OOB_MN
-    static_assert(!(Clear_OOB_MN && !Clear_OOB_K));
-    #pragma unroll
-    for (int m = 0; m < size<1>(S); ++m) {
-        if (Is_even_MN || get<0>(identity_MN(0, m, 0)) < max_MN) {
-            #pragma unroll
-            for (int k = 0; k < size<2>(S); ++k) {
-                if (Is_even_K || predicate_K(k)) {
-                    cute::copy(tiled_copy, S(_, m, k), D(_, m, k));
-                } else if (Clear_OOB_K) {
-                    cute::clear(D(_, m, k));
-                }
-            }
-        } else if (Clear_OOB_MN) {
-            cute::clear(D(_, m, _));
-        }
-    }
-}
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-}  // namespace flash