diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json index 4bd46d2befa..3458be4c59e 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json @@ -1672,6 +1672,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'" diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py index d933eb045c1..3a8cc2ade14 100644 --- a/aiter/ops/mha.py +++ b/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]: ... diff --git a/csrc/cpp_itfs/mha_fwd.cu b/csrc/cpp_itfs/mha_fwd.cu index 10f6aac796f..9c442c618b1 100644 --- a/csrc/cpp_itfs/mha_fwd.cu +++ b/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 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 diff --git a/csrc/include/mha_fwd.h b/csrc/include/mha_fwd.h index d191ab652d1..5e153f2dc5d 100644 --- a/csrc/include/mha_fwd.h +++ b/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 diff --git a/csrc/include/rocm_ops.hpp b/csrc/include/rocm_ops.hpp index 1075a91c349..7cab70338ac 100644 --- a/csrc/include/rocm_ops.hpp +++ b/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", \ diff --git a/csrc/include/torch/mha_v3_varlen_fwd.h b/csrc/include/torch/mha_v3_varlen_fwd.h index d30183dc1e3..0f9d50050c8 100644 --- a/csrc/include/torch/mha_v3_varlen_fwd.h +++ b/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 namespace aiter { @@ -33,6 +33,7 @@ fmha_v3_varlen_fwd(at::Tensor& q, // [total_q, h std::optional v_descale, // [1] or [b, h_k] std::optional gen, std::optional cu_seqlens_q_padded, // [b+1] - std::optional cu_seqlens_k_padded); // [b+1]); + std::optional cu_seqlens_k_padded, // [b+1] + int num_splits = 0); } // namespace torch_itfs } // namespace aiter diff --git a/csrc/kernels/fmha_fwd_v3_splitkv_combine.cu b/csrc/kernels/fmha_fwd_v3_splitkv_combine.cu new file mode 100644 index 00000000000..7a2c3c86064 --- /dev/null +++ b/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 + +#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 +__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(row_count) * kValueDim; + const int64_t lse_plane = static_cast(num_heads) * seqlen_q; + const int64_t lse_index = static_cast(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(row) * kValueDim; + const int64_t out_index = static_cast(token) * out_token_stride + + static_cast(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(partial_out[split * partial_plane + partial_index + value]); + } + out[out_index + value] = hip_bfloat16(result); + } +} + +template +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), + 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(partial_out); + const auto* partial_lse_ptr = static_cast(partial_lse); + auto* out_ptr = static_cast(out); + auto* lse_ptr = static_cast(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 diff --git a/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu b/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu index 57fb9f58b6c..2703d0da4dc 100644 --- a/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu +++ b/csrc/py_itfs_cu/asm_mha_varlen_fwd.cu @@ -8,7 +8,35 @@ #include "mha_fwd.h" namespace aiter { +namespace { + +constexpr int kHd192SplitKvDefaultSplits = 3; +constexpr int kHd192SplitKvQTile = 128; +constexpr int kHd192SplitKvTile = 32; +constexpr int kHd192SplitKvAutoMinKvTiles = 256; + +// 0: auto. 1: force the unsplit production kernel. 2-8: force that split count. +// Auto uses split-3 when KV has at least 256 full 32-token tiles (Sk >= 8192) +// and Q occupancy is not already past 2x CU count. +int select_hd192_splitkv_num_splits( + int num_splits, bool kernel_compatible, int seqlen_q, int nhead, int seqlen_k) +{ + if(num_splits >= kFmhaHd192SplitKvMinSplits) + return num_splits; + if(num_splits == 1 || !kernel_compatible) + return 1; + if(seqlen_k < kHd192SplitKvAutoMinKvTiles * kHd192SplitKvTile) + return 1; + const int q_wgs = ((seqlen_q + kHd192SplitKvQTile - 1) / kHd192SplitKvQTile) * nhead; + if(q_wgs > static_cast(2 * get_num_cu_func())) + return 1; + return kHd192SplitKvDefaultSplits; +} + +} // namespace + namespace torch_itfs { + mha_fwd_args get_asm_mha_varlen_fwd_args(bool has_lse, bool has_dropout_randval, const mask_info &mask, @@ -239,7 +267,7 @@ mha_fwd_args get_asm_mha_varlen_fwd_args(bool has_lse, std::vector -fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] +static fmha_v3_varlen_fwd_impl(at::Tensor &q, // [total_q, hq, d] const at::Tensor &k, // [total_k, hk, d] const at::Tensor &v, // [total_k, hk, d] const at::Tensor &cu_seqlens_q, // [b+1] @@ -266,7 +294,8 @@ fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] std::optional v_descale_, // [1] or [b, h_k] std::optional gen_, std::optional cu_seqlens_q_padded, // [b+1] - std::optional cu_seqlens_k_padded) // [b+1]) + std::optional cu_seqlens_k_padded, // [b+1] + int num_splits) { auto q_dtype = q.dtype(); bool is_qkv_fp8 = q_dtype == at::ScalarType::Float8_e4m3fn || q_dtype == at::ScalarType::Float8_e4m3fnuz; @@ -449,14 +478,58 @@ fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] aiter::ParsePhiloxCudaState, dim3(1), dim3(64), 0, stream, philox_args, rng_state_ptr); } std::optional seqlens_k = std::nullopt; + TORCH_CHECK(num_splits >= 0 && num_splits <= kFmhaHd192SplitKvMaxSplits, + "num_splits must be between 0 and ", + kFmhaHd192SplitKvMaxSplits); + const bool split_forced = num_splits >= kFmhaHd192SplitKvMinSplits; + const bool splitkv_compatible = + get_gpu_arch() == "gfx942" && + q_dtype == torch::kBFloat16 && !is_mi308_device() && + batch_size == 1 && max_seqlen_q == total_q && max_seqlen_k == k.size(0) && + num_heads == num_heads_k && + head_size_q == 192 && head_size_v == 128 && + !is_causal && window_size_left == -1 && window_size_right == -1 && + p_dropout == 0.0f && logits_soft_cap == 0.0f && + bias_type == bias_enum::no_bias && !paged_KV && + !q_descale_.has_value() && !return_dropout_randval && + !cu_seqlens_q_padded.has_value() && !cu_seqlens_k_padded.has_value() && + how_v3_bf16_cvt == 1; + TORCH_CHECK(!split_forced || splitkv_compatible, + "forced gfx942 hd192 split-KV request is incompatible with this input"); + num_splits = select_hd192_splitkv_num_splits( + num_splits, splitkv_compatible, max_seqlen_q, num_heads, max_seqlen_k); + const bool use_hd192_splitkv = num_splits > 1; + if(use_hd192_splitkv && max_seqlen_k > 0) + { + const int kv_tiles = (max_seqlen_k + kHd192SplitKvTile - 1) / kHd192SplitKvTile; + const int split_tiles = (kv_tiles + num_splits - 1) / num_splits; + TORCH_CHECK(split_tiles * (num_splits - 1) < kv_tiles, + "num_splits=", + num_splits, + " would create an empty final KV partition for seqlen_k=", + max_seqlen_k); + } if (max_seqlen_k > 0) { ck_tile::stream_config stream_config{stream}; + at::Tensor o_parts; + at::Tensor lse_parts; + at::Tensor producer_out = out; + at::Tensor producer_lse = softmax_lse; + if(use_hd192_splitkv) + { + o_parts = torch::empty( + {num_splits, total_q, num_heads, head_size_v}, opts.dtype(out_type)); + lse_parts = torch::empty( + {num_splits, num_heads, total_q}, opts.dtype(torch::kFloat32)); + producer_out = o_parts.select(0, 0); + producer_lse = lse_parts.select(0, 0); + } auto drop_seed_offset = std::make_pair(rng_state_ptr, rng_state_ptr + 1); auto args = get_asm_mha_varlen_fwd_args( - has_lse, + has_lse || use_hd192_splitkv, return_dropout_randval, mask, batch_size, @@ -479,8 +552,8 @@ fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] q_descale_, k_descale_, v_descale_, - out, - softmax_lse, + producer_out, + producer_lse, p, softmax_scale, logits_soft_cap, @@ -490,7 +563,27 @@ fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] bias_type, how_v3_bf16_cvt); - float t = aiter::mha_fwd(args, stream_config); + float t; + if(use_hd192_splitkv) + { + t = aiter::fmha_fwd_v3_splitkv(args, num_splits, stream_config); + TORCH_CHECK(t >= 0, "invalid argument for fmha_v3_varlen_fwd"); + aiter::launch_fmha_fwd_v3_splitkv_combine( + o_parts.data_ptr(), + lse_parts.data_ptr(), + out.data_ptr(), + softmax_lse.numel() == 0 ? nullptr : softmax_lse.data_ptr(), + total_q, + num_heads, + num_splits, + out.stride(0), + out.stride(1), + stream); + } + else + { + t = aiter::mha_fwd(args, stream_config); + } TORCH_CHECK(t >= 0, "invalid argument for fmha_v3_varlen_fwd"); } else { @@ -502,5 +595,67 @@ fmha_v3_varlen_fwd(at::Tensor &q, // [total_q, hq, d] return {out, softmax_lse, p, rng_state}; } +std::vector +fmha_v3_varlen_fwd(at::Tensor &q, + const at::Tensor &k, + const at::Tensor &v, + const at::Tensor &cu_seqlens_q, + const at::Tensor &cu_seqlens_k, + int max_seqlen_q, + int max_seqlen_k, + int min_seqlen_q, + float p_dropout, + float softmax_scale, + float logits_soft_cap, + bool zero_tensors, + bool is_causal, + int window_size_left, + int window_size_right, + bool return_softmax_lse, + bool return_dropout_randval, + int how_v3_bf16_cvt, + std::optional out, + std::optional block_table, + std::optional bias, + std::optional alibi_slopes, + std::optional q_descale, + std::optional k_descale, + std::optional v_descale, + std::optional gen, + std::optional cu_seqlens_q_padded, + std::optional cu_seqlens_k_padded, + int num_splits) +{ + return fmha_v3_varlen_fwd_impl(q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + min_seqlen_q, + p_dropout, + softmax_scale, + logits_soft_cap, + zero_tensors, + is_causal, + window_size_left, + window_size_right, + return_softmax_lse, + return_dropout_randval, + how_v3_bf16_cvt, + out, + block_table, + bias, + alibi_slopes, + q_descale, + k_descale, + v_descale, + gen, + cu_seqlens_q_padded, + cu_seqlens_k_padded, + num_splits); +} + } // namespace torch_itfs } // namespace aiter diff --git a/hsa/gfx942/fmha_v3_fwd/MI300/fwd_hd192x128_bf16_rtna_group_splitkv.co b/hsa/gfx942/fmha_v3_fwd/MI300/fwd_hd192x128_bf16_rtna_group_splitkv.co new file mode 100755 index 00000000000..092f8eccf0f Binary files /dev/null and b/hsa/gfx942/fmha_v3_fwd/MI300/fwd_hd192x128_bf16_rtna_group_splitkv.co differ diff --git a/hsa/gfx942/fmha_v3_fwd/fmha_splitkv.csv b/hsa/gfx942/fmha_v3_fwd/fmha_splitkv.csv new file mode 100644 index 00000000000..305375584c1 --- /dev/null +++ b/hsa/gfx942/fmha_v3_fwd/fmha_splitkv.csv @@ -0,0 +1,2 @@ +min_splits,max_splits,knl_name,co_name +2,8,_ZN5aiter42fmha_fwd_hd192x128_bf16_rtna_group_splitkvE,fwd_hd192x128_bf16_rtna_group_splitkv.co diff --git a/op_tests/op_benchmarks/bench_fmha_gfx942_asm_splitkv.py b/op_tests/op_benchmarks/bench_fmha_gfx942_asm_splitkv.py new file mode 100755 index 00000000000..99635fcca7c --- /dev/null +++ b/op_tests/op_benchmarks/bench_fmha_gfx942_asm_splitkv.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Matched wall-clock benchmark for production ASM versus split-KV ASM.""" + +import triton # noqa: F401 # isort: skip # Must precede torch on this ROCm environment. + +import argparse +import math +import statistics +import time + +import torch + +from aiter.ops.mha import fmha_v3_varlen_fwd + + +def production_asm(q, k, v, cu_q, cu_k, scale, num_splits): + out = torch.empty(q.shape[0], q.shape[1], 128, dtype=q.dtype, device=q.device) + result, _, _, _ = fmha_v3_varlen_fwd( + q, + k, + v, + cu_q, + cu_k, + q.shape[0], + k.shape[0], + 0, + 0.0, + scale, + 0.0, + False, + False, + -1, + -1, + False, + False, + 1, + out=out, + num_splits=num_splits, + ) + return result + + +def measure_paired(fn, num_splits: int, warmup: int, samples: int): + for _ in range(warmup): + fn(1) + fn(num_splits) + torch.cuda.synchronize() + + def timed(splits: int) -> float: + start = time.perf_counter() + fn(splits) + torch.cuda.synchronize() + return (time.perf_counter() - start) * 1e3 + + baseline_values = [] + split_values = [] + for sample in range(samples): + if sample % 2 == 0: + baseline_values.append(timed(1)) + split_values.append(timed(num_splits)) + else: + split_values.append(timed(num_splits)) + baseline_values.append(timed(1)) + return baseline_values, split_values + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--sq", type=int, default=4096) + parser.add_argument("--sk", type=int, default=42700) + parser.add_argument("--heads", type=int, default=12) + parser.add_argument("--splits", type=int, default=3) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--samples", type=int, default=30) + args = parser.parse_args() + + generator = torch.Generator(device="cpu").manual_seed(0) + q = torch.randn( + args.sq, args.heads, 192, dtype=torch.bfloat16, generator=generator + ).cuda() + k = torch.randn( + args.sk, args.heads, 192, dtype=torch.bfloat16, generator=generator + ).cuda() + v = torch.randn( + args.sk, args.heads, 128, dtype=torch.bfloat16, generator=generator + ).cuda() + cu_q = torch.tensor([0, args.sq], dtype=torch.int32, device="cuda") + cu_k = torch.tensor([0, args.sk], dtype=torch.int32, device="cuda") + scale = 1.0 / math.sqrt(192) + + if not 2 <= args.splits <= 8: + raise ValueError("splits must be between 2 and 8") + run = lambda num_splits: production_asm(q, k, v, cu_q, cu_k, scale, num_splits) + reference = run(1) + actual = run(args.splits) + if not torch.isfinite(reference).all() or not torch.isfinite(actual).all(): + raise RuntimeError("correctness gate failed: kernel produced non-finite output") + denom = (reference.double().square().sum() + actual.double().square().sum()).item() + if not math.isfinite(denom) or denom <= 0.0: + raise RuntimeError( + f"correctness gate failed: cosine denominator is not finite and positive ({denom})" + ) + cosine_difference = ( + 1.0 - 2.0 * (reference.double() * actual.double()).sum().item() / denom + ) + if not math.isfinite(cosine_difference) or cosine_difference >= 1e-4: + raise RuntimeError( + f"correctness gate failed: cosine difference={cosine_difference}" + ) + + baseline_ms, split_ms = measure_paired(run, args.splits, args.warmup, args.samples) + flop = args.heads * 2 * args.sq * args.sk * (192 + 128) + baseline_median = statistics.median(baseline_ms) + split_median = statistics.median(split_ms) + print(f"gfx942 Sq={args.sq} Sk={args.sk} H={args.heads} S={args.splits}") + print( + "production_asm " + f"median={baseline_median:.3f}ms min={min(baseline_ms):.3f}ms " + f"max={max(baseline_ms):.3f}ms tflops={flop / baseline_median / 1e9:.1f}" + ) + print( + "splitkv_asm " + f"median={split_median:.3f}ms min={min(split_ms):.3f}ms " + f"max={max(split_ms):.3f}ms tflops={flop / split_median / 1e9:.1f}" + ) + print( + f"speedup={(baseline_median / split_median):.3f}x " + f"latency_reduction={(1.0 - split_median / baseline_median) * 100:.1f}% " + f"cosine_difference={cosine_difference:.3e}" + ) + + +if __name__ == "__main__": + main() diff --git a/op_tests/test_fmha_gfx942_asm_splitkv.py b/op_tests/test_fmha_gfx942_asm_splitkv.py new file mode 100644 index 00000000000..7416c770e1b --- /dev/null +++ b/op_tests/test_fmha_gfx942_asm_splitkv.py @@ -0,0 +1,478 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""Correctness + performance test for gfx942 packed-varlen hd192 split-KV FMHA. + +Public API: aiter.flash_attn_varlen_func (the path vLLM / the ticket calls) +Ops layer: aiter.ops.mha.fmha_v3_varlen_fwd (num_splits last; default 0) + +Built to the aiter op-test standard (see .claude/skills/aiter-op-test). + +num_splits (asm_mha_varlen_fwd.cu): 0 = auto, 1 = unsplit production kernel, +2-8 = forced split-KV. Auto uses split-3 when the kernel contract matches, +Sk >= 8192, and Q_tiles * heads <= 2 * CU. + +q/k/v are packed THD, batch 1, BF16, D_QK=192 / D_V=128, non-causal — the +layout the model uses for this kernel. +""" + +import triton # noqa: F401 # isort: skip # Must precede torch on this ROCm environment. + +import argparse +import itertools +import math + +import pandas as pd +import torch + +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_cu_num, get_device_name, get_gfx +from aiter.ops.mha import flash_attn_varlen_func, fmha_v3_varlen_fwd +from aiter.test_common import assertAllclose, benchmark, run_perftest + +torch.set_default_device("cuda") + +# fmha_fwd_v3_splitkv requires gfx942 and rejects MI308. The public wrapper +# still runs on other gfx942 cards; this file only times the split-KV kernel. +SUPPORTED_GFX = ["gfx942"] + +HD_QK = 192 +HD_V = 128 +KV_TILE = 32 + +# (sq, sk, hq) — packed batch-1. Ticket shape is last. +_SHAPES = [ + (129, 511, 4), # Q/KV tile tails + (129, 2048, 4), # split-count coverage + (4096, 8191, 12), # last length that stays unsplit + (4096, 8192, 12), # first auto split-3 + (4096, 42700, 12), # Kimi-K3 packed-varlen ticket +] + + +def run_torch(q, k, v, scale): + """Packed THD reference, fp32 math. Not timed, not in the table. + + Returns (out in q.dtype, lse in fp32 with shape [H, Sq]). + Empty K matches the C++ empty-sequence branch: zeros and +inf LSE. + Runs on CPU in Q tiles so the softmax reduction fits in host memory. + """ + sq, hq, _ = q.shape + sk = k.shape[0] + if sk == 0: + out = torch.zeros(sq, hq, HD_V, dtype=q.dtype, device=q.device) + lse = torch.full((hq, sq), float("inf"), dtype=torch.float32, device=q.device) + return out, lse + q_cpu = q.float().cpu() + k_h = k.float().cpu().permute(1, 2, 0).contiguous() + v_h = v.float().cpu().permute(1, 0, 2).contiguous() + out = torch.empty(sq, hq, HD_V, dtype=torch.float32, device="cpu") + lse = torch.empty(hq, sq, dtype=torch.float32, device="cpu") + q_tile = 128 + for q0 in range(0, sq, q_tile): + q1 = min(q0 + q_tile, sq) + q_h = q_cpu[q0:q1].permute(1, 0, 2).contiguous() + scores = torch.bmm(q_h, k_h) * scale + lse[:, q0:q1] = torch.logsumexp(scores, dim=-1) + out[q0:q1] = torch.bmm(torch.softmax(scores, dim=-1), v_h).permute(1, 0, 2) + return out.to(device=q.device, dtype=q.dtype), lse.to(device=q.device) + + +def _flops_bytes(hq, sq, sk, esz): + """Attention roofline: 2 GEMMs (QK^T, PV), HBM traffic q+k+v+o.""" + flops = 2.0 * hq * sq * sk * (HD_QK + HD_V) + nbytes = (sq * hq * HD_QK + sk * hq * HD_QK + sk * hq * HD_V + sq * hq * HD_V) * esz + return flops, nbytes + + +def _empty_final_split(sk, num_splits): + kv_tiles = (sk + KV_TILE - 1) // KV_TILE + split_tiles = (kv_tiles + num_splits - 1) // num_splits + return split_tiles * (num_splits - 1) >= kv_tiles + + +def _inputs(sq, sk, hq): + q = torch.randn(sq, hq, HD_QK, dtype=dtypes.bf16) + k = ( + torch.empty(0, hq, HD_QK, dtype=dtypes.bf16) + if sk == 0 + else torch.randn(sk, hq, HD_QK, dtype=dtypes.bf16) + ) + v = ( + torch.empty(0, hq, HD_V, dtype=dtypes.bf16) + if sk == 0 + else torch.randn(sk, hq, HD_V, dtype=dtypes.bf16) + ) + cu_q = torch.tensor([0, sq], dtype=torch.int32) + cu_k = torch.tensor([0, sk], dtype=torch.int32) + scale = 1.0 / math.sqrt(HD_QK) + return q, k, v, cu_q, cu_k, scale + + +def _v3_fwd(q, k, v, cu_q, cu_k, scale, num_splits, return_lse, out=None, causal=False): + """Full fmha_v3_varlen_fwd entry: caller args reach the impl, including out=.""" + if out is None: + out = torch.empty(q.shape[0], q.shape[1], HD_V, dtype=q.dtype, device=q.device) + return fmha_v3_varlen_fwd( + q, + k, + v, + cu_q, + cu_k, + q.shape[0], + k.shape[0], + 0, + 0.0, + scale, + 0.0, + False, + causal, + -1, + -1, + return_lse, + False, + 1, + out=out, + num_splits=num_splits, + ) + + +def _public(q, k, v, cu_q, cu_k, scale, return_lse, out=None): + # Preallocated out= is the buffer the model can pass through the public API. + if out is None: + out = torch.empty(q.shape[0], q.shape[1], HD_V, dtype=q.dtype, device=q.device) + result = flash_attn_varlen_func( + q, + k, + v, + cu_q, + cu_k, + q.shape[0], + k.shape[0], + softmax_scale=scale, + causal=False, + return_lse=return_lse, + out=out, + ) + if return_lse: + return result[0], result[1] + return result, None + + +@benchmark() +def test_fmha_gfx942_asm_splitkv(sq, sk, hq, num_splits, return_lse): + torch.manual_seed(0) + q = torch.randn(sq, hq, HD_QK, dtype=dtypes.bf16) + k = torch.randn(sk, hq, HD_QK, dtype=dtypes.bf16) + v = torch.randn(sk, hq, HD_V, dtype=dtypes.bf16) + cu_q = torch.tensor([0, sq], dtype=torch.int32) + cu_k = torch.tensor([0, sk], dtype=torch.int32) + scale = 1.0 / math.sqrt(HD_QK) + out_buf = torch.empty(sq, hq, HD_V, dtype=q.dtype) + + ref_out, ref_lse = run_torch(q, k, v, scale) + flops, nbytes = _flops_bytes(hq, sq, sk, q.element_size()) + + candidates = { + # Production unsplit ASM (num_splits=1). + "unsplit": lambda: _v3_fwd(q, k, v, cu_q, cu_k, scale, 1, return_lse)[:2], + # The path the model actually runs. + "public": lambda: _public(q, k, v, cu_q, cu_k, scale, return_lse, out=out_buf), + } + # Forced split is rejected when the last KV partition would be empty. + if not _empty_final_split(sk, num_splits): + candidates["splitkv"] = lambda: _v3_fwd( + q, k, v, cu_q, cu_k, scale, num_splits, return_lse + )[:2] + + ret = { + "gfx": get_gfx(), + "cu": get_cu_num(), + "q_wgs": ((sq + 127) // 128) * hq, + } + outs = {} + for name, fn in candidates.items(): + (out, lse), us = run_perftest(fn, num_rotate_args=1) + outs[name] = out + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 if us else float("nan") + ret[f"{name} TB/s"] = nbytes / us / 1e6 if us else float("nan") + ret[f"{name} err"] = assertAllclose( + ref_out.to(dtypes.fp32), + out.to(dtypes.fp32), + rtol=2e-2, + atol=2e-2, + msg=f"{name} O sq={sq} sk={sk} hq={hq} ns={num_splits}", + ) + if return_lse: + assertAllclose( + ref_lse.to(dtypes.fp32), + lse.to(dtypes.fp32), + rtol=2e-2, + atol=2e-2, + msg=f"{name} LSE sq={sq} sk={sk} hq={hq} ns={num_splits}", + ) + auto_split = sk >= 8192 and ret["q_wgs"] <= 2 * ret["cu"] + if "public" in outs and "unsplit" in outs: + public_eq_unsplit = torch.equal(outs["public"], outs["unsplit"]) + ret["public_eq_unsplit"] = int(public_eq_unsplit) + if auto_split: + assert not public_eq_unsplit, ( + sq, + sk, + hq, + "public auto should not be unsplit", + ) + else: + assert public_eq_unsplit, (sq, sk, hq, "public auto should stay unsplit") + if "public" in outs and "splitkv" in outs: + public_eq_splitkv = torch.equal(outs["public"], outs["splitkv"]) + ret["public_eq_splitkv"] = int(public_eq_splitkv) + if auto_split and num_splits == 3: + assert public_eq_splitkv, (sq, sk, hq, "public auto should be split-3") + return ret + + +@benchmark() +def test_fmha_gfx942_asm_splitkv_empty_k(sq, hq, num_splits): + torch.manual_seed(1) + q = torch.randn(sq, hq, HD_QK, dtype=dtypes.bf16) + k = torch.empty(0, hq, HD_QK, dtype=dtypes.bf16) + v = torch.empty(0, hq, HD_V, dtype=dtypes.bf16) + cu_q = torch.tensor([0, sq], dtype=torch.int32) + cu_k = torch.tensor([0, 0], dtype=torch.int32) + scale = 1.0 / math.sqrt(HD_QK) + ref_out, _ = run_torch(q, k, v, scale) + flops, nbytes = _flops_bytes(hq, sq, 0, q.element_size()) + + candidates = { + "unsplit": lambda: _v3_fwd(q, k, v, cu_q, cu_k, scale, 1, False)[:2], + "splitkv": lambda: _v3_fwd(q, k, v, cu_q, cu_k, scale, num_splits, False)[:2], + "public": lambda: _public(q, k, v, cu_q, cu_k, scale, False), + } + ret = {"gfx": get_gfx()} + for name, fn in candidates.items(): + (out, _), us = run_perftest(fn, num_rotate_args=1) + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 if us else float("nan") + ret[f"{name} TB/s"] = nbytes / us / 1e6 if us else float("nan") + ret[f"{name} err"] = assertAllclose( + ref_out.to(dtypes.fp32), + out.to(dtypes.fp32), + rtol=0, + atol=0, + msg=f"{name} empty-K sq={sq} hq={hq} ns={num_splits}", + ) + return ret + + +def _check_empty_partition_rejected(): + q, k, v, cu_q, cu_k, scale = _inputs(129, 511, 4) + try: + _v3_fwd(q, k, v, cu_q, cu_k, scale, 5, False) + except RuntimeError as err: + if "empty final KV partition" not in str(err): + raise + return + raise AssertionError("forced split-5 on Sk=511 should reject an empty partition") + + +def _check_compile_outputs(): + q, k, v, cu_q, cu_k, scale = _inputs(129, 2048, 4) + sq, sk = q.shape[0], k.shape[0] + + def call(q, k, v): + return fmha_v3_varlen_fwd( + q, + k, + v, + cu_q, + cu_k, + sq, + sk, + 0, + 0.0, + scale, + 0.0, + False, + False, + -1, + -1, + True, + False, + 1, + num_splits=3, + ) + + eager = call(q, k, v) + compiled = torch.compile(call, fullgraph=True)(q, k, v) + for idx, (eager_t, compiled_t) in enumerate(zip(eager, compiled)): + assert eager_t.dtype == compiled_t.dtype, (idx, eager_t.dtype, compiled_t.dtype) + assert eager_t.shape == compiled_t.shape, (idx, eager_t.shape, compiled_t.shape) + assert eager[0].dtype == q.dtype + assert eager[1].dtype == torch.float32 + assert eager[2].dtype == q.dtype + assert eager[3].dtype == torch.int64 + assertAllclose( + eager[0].to(dtypes.fp32), + compiled[0].to(dtypes.fp32), + rtol=2e-2, + atol=2e-2, + msg="torch.compile O", + ) + assertAllclose( + eager[1].to(dtypes.fp32), + compiled[1].to(dtypes.fp32), + rtol=2e-4, + atol=2e-4, + msg="torch.compile LSE", + ) + + +def _check_cuda_graph(): + # sk=8192 is the first auto-select length; this captures split producer + combine. + q, k, v, cu_q, cu_k, scale = _inputs(129, 8192, 4) + sq, sk = q.shape[0], k.shape[0] + for _ in range(3): + flash_attn_varlen_func( + q, k, v, cu_q, cu_k, sq, sk, softmax_scale=scale, causal=False + ) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = flash_attn_varlen_func( + q, k, v, cu_q, cu_k, sq, sk, softmax_scale=scale, causal=False + ) + graph.replay() + first = captured.clone() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(captured, first) + + +def _check_forced_split_writes_out(): + q, k, v, cu_q, cu_k, scale = _inputs(129, 2048, 4) + out = torch.empty(q.shape[0], q.shape[1], HD_V, dtype=q.dtype) + sentinel = out.data_ptr() + result = _v3_fwd(q, k, v, cu_q, cu_k, scale, 3, False, out=out) + assert result[0].data_ptr() == sentinel, "forced split must write the caller out=" + assert torch.equal(result[0], out) + ref_out, _ = run_torch(q, k, v, scale) + assertAllclose( + ref_out.to(dtypes.fp32), + out.to(dtypes.fp32), + rtol=2e-2, + atol=2e-2, + msg="forced split out=", + ) + + +def _check_forced_split_causal_rejected(): + q, k, v, cu_q, cu_k, scale = _inputs(129, 2048, 4) + try: + _v3_fwd(q, k, v, cu_q, cu_k, scale, 3, False, causal=True) + except RuntimeError as err: + if "incompatible with this input" not in str(err): + raise + return + raise AssertionError("forced split with causal=True should raise") + + +def _check_non_lse(): + q, k, v, cu_q, cu_k, scale = _inputs(129, 2048, 4) + ref_out, _ = run_torch(q, k, v, scale) + out, lse, _, _ = _v3_fwd(q, k, v, cu_q, cu_k, scale, 3, False) + assert lse.numel() == 0 + assertAllclose( + ref_out.to(dtypes.fp32), + out.to(dtypes.fp32), + rtol=2e-2, + atol=2e-2, + msg="forced split return_lse=False", + ) + + +def main(): + if get_gfx() not in SUPPORTED_GFX or get_device_name() == "MI308": + aiter.logger.warning( + "gfx942 hd192 split-KV unsupported on %s; skipping", + get_device_name() if get_gfx() == "gfx942" else get_gfx(), + ) + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="config input of test", + ) + parser.add_argument( + "-d", + "--dtype", + type=dtypes.str2Dtype, + nargs="*", + default=[dtypes.bf16], + help="""Data type. + e.g.: -d bf16""", + ) + parser.add_argument( + "-s", + "--shapes", + type=dtypes.str2tuple, + nargs="*", + default=_SHAPES, + help="shape(s) as sq,sk,hq (default: tails, auto-select bounds, ticket)", + ) + parser.add_argument( + "-ns", + "--num_splits", + type=int, + nargs="*", + default=list(range(2, 9)), + help="forced split count(s) for the splitkv candidate (default: 2..8)", + ) + parser.add_argument( + "--lse", + type=int, + nargs="*", + choices=[0, 1], + default=[1], + help="return_lse: 0=inference 1=training (default: 1)", + ) + args = parser.parse_args() + + for dtype in args.dtype: + if dtype != dtypes.bf16: + aiter.logger.warning("hd192 split-KV is bf16-only; skipping %s", dtype) + continue + df = [] + for shape, num_splits, return_lse in itertools.product( + args.shapes, args.num_splits, args.lse + ): + sq, sk, hq = shape + df.append( + test_fmha_gfx942_asm_splitkv(sq, sk, hq, num_splits, bool(return_lse)) + ) + df = pd.DataFrame(df) + aiter.logger.info( + "fmha_gfx942_asm_splitkv summary (markdown):\n%s", + df.to_markdown(index=False), + ) + + empty_rows = [ + test_fmha_gfx942_asm_splitkv_empty_k(129, 4, num_splits) + for num_splits in args.num_splits + ] + aiter.logger.info( + "fmha_gfx942_asm_splitkv empty-K summary (markdown):\n%s", + pd.DataFrame(empty_rows).to_markdown(index=False), + ) + + _check_empty_partition_rejected() + _check_forced_split_writes_out() + _check_forced_split_causal_rejected() + _check_non_lse() + _check_compile_outputs() + _check_cuda_graph() + + +if __name__ == "__main__": + main()