[ASM] [HIP] [JIT] [gfx942] Optimize long-context FMHA with split-KV

aiter/jit/optCompilerConfig.json

@@ -1643,6 +1643,7 @@
     "module_fmha_v3_varlen_fwd": {
         "srcs": [
             "f'{AITER_CSRC_DIR}/kernels/mha_common.cu'",
+            "f'{AITER_CSRC_DIR}/kernels/fmha_fwd_v3_splitkv_combine.cu'",
             "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_mha_varlen_fwd.cu'",
             "f'{AITER_CSRC_DIR}/pybind/mha_varlen_fwd_asm_pybind.cu'",
             "f'{AITER_CSRC_DIR}/cpp_itfs/mha_fwd.cu'"

aiter/ops/mha.py

@@ -1089,6 +1089,7 @@ def gen_fmha_v3_varlen_fwd_fake_tensor(
     gen: torch.Generator | None = None,
     cu_seqlens_q_padded: torch.Tensor | None = None,
     cu_seqlens_k_padded: torch.Tensor | None = None,
+    num_splits: int = 0,
 ) -> tuple[Tensor, Tensor, Tensor, Tensor]:
     device = q.device
     dtype = q.dtype
@@ -1115,7 +1116,7 @@ def gen_fmha_v3_varlen_fwd_fake_tensor(
         p_shape = (num_heads, total_q, max_seqlen_k)
         p_tensor = torch.empty(p_shape, device=device, dtype=torch.uint8)
     else:
-        p_tensor = torch.empty((0,), device=device)
+        p_tensor = torch.empty((0,), device=device, dtype=dtype)
 
     rng_state_tensor = torch.empty((2,), device=device, dtype=torch.int64)
 
@@ -1156,6 +1157,7 @@ def fmha_v3_varlen_fwd(
     gen: torch.Generator | None = None,
     cu_seqlens_q_padded: torch.Tensor | None = None,
     cu_seqlens_k_padded: torch.Tensor | None = None,
+    num_splits: int = 0,
 ) -> tuple[Tensor, Tensor, Tensor, Tensor]: ...
 
 

csrc/cpp_itfs/mha_fwd.cu

@@ -282,6 +282,70 @@ float fmha_fwd_v3(mha_fwd_args a, const ck_tile::stream_config& s)
         impl_ptr->launch_kernel({args_ptr, arg_size_ptr, gdx, gdy, gdz, bdx, 1, 1, s_.stream_id_});
     });
 }
+
+float fmha_fwd_v3_splitkv(mha_fwd_args a,
+                          int num_splits,
+                          const ck_tile::stream_config& stream_config)
+{
+    const std::string arch = get_gpu_arch();
+    AITER_CHECK(arch == "gfx942" && !is_mi308_device(),
+                "fmha_fwd_v3_splitkv requires a gfx942 MI300-family device");
+    AITER_CHECK(a.data_type == "bf16" && a.is_group_mode && a.batch == 1 && a.hdim_q == 192 &&
+                    a.hdim_v == 128 && a.nhead_q == a.nhead_k && a.mask_type == 0 &&
+                    a.bias_type == 0 && a.p_drop == 0.0f &&
+                    num_splits >= kFmhaHd192SplitKvMinSplits &&
+                    num_splits <= kFmhaHd192SplitKvMaxSplits,
+                "invalid argument for fmha_fwd_v3_splitkv");
+
+    const fmha_v3_fwdConfig* selected = nullptr;
+    for(const auto& entry : cfg_fmha_splitkv)
+    {
+        const auto& cfg = entry.second;
+        if(cfg.arch == arch && num_splits >= cfg.min_splits && num_splits <= cfg.max_splits)
+        {
+            selected = &cfg;
+            break;
+        }
+    }
+    AITER_CHECK(selected != nullptr, "no matching fmha_fwd_v3_splitkv code object");
+
+    fmha_fwd_v3_args args;
+    init_fmha_fwd_v3_args(args, a, 128, arch);
+    const int kv_tiles    = (a.seqlen_k + 31) / 32;
+    const int split_tiles = (kv_tiles + num_splits - 1) / num_splits;
+    AITER_CHECK(split_tiles * (num_splits - 1) < kv_tiles,
+                "split partition would contain an empty final split");
+
+    // This variant maps grid Y to the split ID, s_v_Bs to keys per split,
+    // and s_o_Bs to the byte stride between split-major output planes.
+    args.s_v_Bs = split_tiles * 32;
+    args.s_o_Bs = a.seqlen_q * args.s_o_Seqs;
+    args.s_lse  = 1;
+
+    static SynchronizedCache<std::string_view, AiterAsmKernel> kernels;
+    const std::string co_name = get_kernel_co_name(selected->co_name, arch);
+    AiterAsmKernel* kernel    = &kernels.get_or_create(selected->knl_name, [&]() {
+        return AiterAsmKernel(selected->knl_name.c_str(), co_name.c_str());
+    });
+
+    size_t arg_size   = sizeof(args);
+    const int q_tiles = (a.seqlen_q + 127) / 128;
+    return ck_tile::launch_kernel(stream_config,
+                                  [=](const ck_tile::stream_config& stream_config_) mutable {
+                                      void* args_ptr       = &args;
+                                      size_t* arg_size_ptr = &arg_size;
+                                      kernel->launch_kernel({args_ptr,
+                                                             arg_size_ptr,
+                                                             a.nhead_q,
+                                                             num_splits,
+                                                             q_tiles,
+                                                             256,
+                                                             1,
+                                                             1,
+                                                             stream_config_.stream_id_});
+                                  });
+}
+
 #endif
 
 #if FAV2_ON

csrc/include/mha_fwd.h

@@ -361,4 +361,24 @@ __attribute__((visibility("default"))) float mha_fwd(mha_fwd_args args,
                                                      const ck_tile::stream_config& s);
 
 float fmha_fwd_v3(mha_fwd_args a, const ck_tile::stream_config& s);
+
+inline constexpr int kFmhaHd192SplitKvMinSplits = 2;
+inline constexpr int kFmhaHd192SplitKvMaxSplits = 8;
+
+float fmha_fwd_v3_splitkv(mha_fwd_args args,
+                          int num_splits,
+                          const ck_tile::stream_config& stream_config);
+
+// Partial layouts are [split, query, head, value] for output and
+// [split, head, query] for natural-log LSE.
+void launch_fmha_fwd_v3_splitkv_combine(const void* partial_out,
+                                        const void* partial_lse,
+                                        void* out,
+                                        void* lse,
+                                        int seqlen_q,
+                                        int num_heads,
+                                        int num_splits,
+                                        int64_t out_token_stride,
+                                        int64_t out_head_stride,
+                                        hipStream_t stream);
 } // namespace aiter

csrc/include/rocm_ops.hpp

@@ -1179,7 +1179,8 @@ namespace py = pybind11;
           py::arg("v_descale")           = std::nullopt, \
           py::arg("gen")                 = std::nullopt, \
           py::arg("cu_seqlens_q_padded") = std::nullopt, \
-          py::arg("cu_seqlens_k_padded") = std::nullopt);
+          py::arg("cu_seqlens_k_padded") = std::nullopt, \
+          py::arg("num_splits")          = 0);
 
 #define MHA_VARLEN_BWD_PYBIND                            \
     m.def("mha_varlen_bwd",                              \

csrc/include/torch/mha_v3_varlen_fwd.h

@@ -1,6 +1,6 @@
 #pragma once
 // SPDX-License-Identifier: MIT
-// Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 #include <torch/extension.h>
 
 namespace aiter {
@@ -33,6 +33,7 @@ fmha_v3_varlen_fwd(at::Tensor& q,                                 // [total_q, h
                    std::optional<const at::Tensor> v_descale,    // [1] or [b, h_k]
                    std::optional<at::Generator> gen,
                    std::optional<const at::Tensor> cu_seqlens_q_padded, // [b+1]
-                   std::optional<const at::Tensor> cu_seqlens_k_padded); // [b+1]);
+                   std::optional<const at::Tensor> cu_seqlens_k_padded, // [b+1]
+                   int num_splits = 0);
 } // namespace torch_itfs
 } // namespace aiter

csrc/kernels/fmha_fwd_v3_splitkv_combine.cu

@@ -0,0 +1,163 @@
+// SPDX-License-Identifier: MIT
+// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <hip/hip_bfloat16.h>
+
+#include "aiter_hip_common.h"
+#include "mha_fwd.h"
+
+namespace aiter {
+namespace {
+
+constexpr int kValueDim     = 128;
+constexpr int kCombineBlock = 256;
+constexpr int kWaveSize     = 64;
+constexpr int kRowsPerBlock = kCombineBlock / kWaveSize;
+
+template <int NumSplits>
+__global__ __launch_bounds__(kCombineBlock) void fmha_fwd_v3_splitkv_combine_kernel(
+    const hip_bfloat16* __restrict__ partial_out,
+    const float* __restrict__ partial_lse,
+    hip_bfloat16* __restrict__ out,
+    float* __restrict__ lse_out,
+    int seqlen_q,
+    int num_heads,
+    int64_t out_token_stride,
+    int64_t out_head_stride)
+{
+    const int lane      = threadIdx.x & (kWaveSize - 1);
+    const int wave      = threadIdx.x / kWaveSize;
+    const int row       = blockIdx.x * kRowsPerBlock + wave;
+    const int row_count = seqlen_q * num_heads;
+    if(row >= row_count)
+        return;
+
+    const int token             = row / num_heads;
+    const int head              = row - token * num_heads;
+    const int64_t partial_plane = static_cast<int64_t>(row_count) * kValueDim;
+    const int64_t lse_plane     = static_cast<int64_t>(num_heads) * seqlen_q;
+    const int64_t lse_index     = static_cast<int64_t>(head) * seqlen_q + token;
+
+    float weights[NumSplits];
+    if(lane == 0)
+    {
+        float lse_max = partial_lse[lse_index];
+#pragma unroll
+        for(int split = 1; split < NumSplits; ++split)
+            lse_max = fmaxf(lse_max, partial_lse[split * lse_plane + lse_index]);
+
+        float denominator = 0.0f;
+#pragma unroll
+        for(int split = 0; split < NumSplits; ++split)
+        {
+            weights[split] = expf(partial_lse[split * lse_plane + lse_index] - lse_max);
+            denominator += weights[split];
+        }
+        denominator         = fmaxf(denominator, 1.0e-20f);
+        const float inverse = 1.0f / denominator;
+#pragma unroll
+        for(int split = 0; split < NumSplits; ++split)
+            weights[split] *= inverse;
+
+        if(lse_out != nullptr)
+            lse_out[lse_index] = lse_max + logf(denominator);
+    }
+
+#pragma unroll
+    for(int split = 0; split < NumSplits; ++split)
+        weights[split] = __shfl(weights[split], 0);
+
+    const int64_t partial_index = static_cast<int64_t>(row) * kValueDim;
+    const int64_t out_index     = static_cast<int64_t>(token) * out_token_stride +
+                              static_cast<int64_t>(head) * out_head_stride;
+    for(int value = lane; value < kValueDim; value += kWaveSize)
+    {
+        float result = 0.0f;
+#pragma unroll
+        for(int split = 0; split < NumSplits; ++split)
+        {
+            result +=
+                weights[split] *
+                static_cast<float>(partial_out[split * partial_plane + partial_index + value]);
+        }
+        out[out_index + value] = hip_bfloat16(result);
+    }
+}
+
+template <int NumSplits>
+void launch_combine(const hip_bfloat16* partial_out,
+                    const float* partial_lse,
+                    hip_bfloat16* out,
+                    float* lse,
+                    int seqlen_q,
+                    int num_heads,
+                    int64_t out_token_stride,
+                    int64_t out_head_stride,
+                    dim3 grid,
+                    hipStream_t stream)
+{
+    hipLaunchKernelGGL((fmha_fwd_v3_splitkv_combine_kernel<NumSplits>),
+                       grid,
+                       dim3(kCombineBlock),
+                       0,
+                       stream,
+                       partial_out,
+                       partial_lse,
+                       out,
+                       lse,
+                       seqlen_q,
+                       num_heads,
+                       out_token_stride,
+                       out_head_stride);
+}
+
+} // namespace
+
+void launch_fmha_fwd_v3_splitkv_combine(const void* partial_out,
+                                        const void* partial_lse,
+                                        void* out,
+                                        void* lse,
+                                        int seqlen_q,
+                                        int num_heads,
+                                        int num_splits,
+                                        int64_t out_token_stride,
+                                        int64_t out_head_stride,
+                                        hipStream_t stream)
+{
+    AITER_CHECK(num_splits >= kFmhaHd192SplitKvMinSplits &&
+                    num_splits <= kFmhaHd192SplitKvMaxSplits,
+                "unsupported fmha_fwd_v3_splitkv combine split count");
+
+    const int row_count = seqlen_q * num_heads;
+    const dim3 grid((row_count + kRowsPerBlock - 1) / kRowsPerBlock);
+    const auto* partial_out_ptr = static_cast<const hip_bfloat16*>(partial_out);
+    const auto* partial_lse_ptr = static_cast<const float*>(partial_lse);
+    auto* out_ptr               = static_cast<hip_bfloat16*>(out);
+    auto* lse_ptr               = static_cast<float*>(lse);
+
+    using combine_launcher                        = decltype(&launch_combine<2>);
+    static constexpr combine_launcher launchers[] = {
+        nullptr,
+        nullptr,
+        &launch_combine<2>,
+        &launch_combine<3>,
+        &launch_combine<4>,
+        &launch_combine<5>,
+        &launch_combine<6>,
+        &launch_combine<7>,
+        &launch_combine<8>,
+    };
+    launchers[num_splits](partial_out_ptr,
+                          partial_lse_ptr,
+                          out_ptr,
+                          lse_ptr,
+                          seqlen_q,
+                          num_heads,
+                          out_token_stride,
+                          out_head_stride,
+                          grid,
+                          stream);
+    HIP_CALL(hipGetLastError());
+}
+
+} // namespace aiter