ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 11:07:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#6526 (commit 3e01710)

Browse Source 
feat: [CK Tile] mxfp8 support for qr async pipeline
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Closes: #6526

## Motivation

Adds MXFP8 support to the QR async FMHA forward pipeline by enabling MX
quant scale paths, loosening padding constraints, and updating
tests/codegen to cover the new combinations.

## Technical Details

Changes:

- Extend BlockFmhaPipelineQRKSVSAsync to support MX quant scales
(Q/K/V/P scale tiles) and adjust padding/alignment behavior accordingly.

- Update FMHA fwd tests to include head-dim adjustment helpers and add a
new “General” parameterized test suite.

- Update codegen pipeline constraints/selection for qr_async, and adjust
tile-window vectorization logic for packed types.

## Test Plan

Running CI with extended set of tests

## Test Result

CI passing

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
This commit is contained in:
Zoltán Lakatos
2026-07-13 10:44:45 +00:00
committed by assistant-librarian[bot]
parent 4975bd0c8e
commit cb859854a7
5 changed files with 330 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
View File

@@ -355,18 +355,12 @@ class FmhaFwdApiTrait:
@property
def dcheck(self) -> str:
if self.pipeline_tag == "qr_async":
vec = int((32 * 4) / DTYPE_BITS[self.dtype])
if self.dpad == "t":
return f"a.hdim_q % {vec} == 0"
else:
assert False
elif self.pipeline_tag == "qr_hpad":
if self.pipeline_tag == "qr_hpad":
if self.dpad == "t":
return "a.hdim_q % 8 == 0"
else:
assert False
elif self.pipeline_tag in ["qr", "qs", "qr_async_trload", "qr_async_trload_v3"]:
elif self.pipeline_tag in ["qr", "qs", "qr_async", "qr_async_trload", "qr_async_trload_v3"]:
bk0submax = K0_MAX_SUBMAX_MAP[self.bk0max]
if self.dpad == "t":
return f"true /*a.hdim_q % {bk0submax} != 0*/" # TODO: order of get_pipelines() matters! (ugly)
@@ -377,18 +371,12 @@ class FmhaFwdApiTrait:
@property
def dvcheck(self) -> str:
if self.pipeline_tag == "qr_async":
vec = int((32 * 4) / DTYPE_BITS[self.dtype])
if self.dvpad == "t":
return f"a.hdim_v % {vec} == 0"
else:
assert False
elif self.pipeline_tag == "qr_hpad":
if self.pipeline_tag == "qr_hpad":
if self.dvpad == "t":
return "a.hdim_v % 8 == 0"
else:
assert False
elif self.pipeline_tag in ["qr", "qs", "qr_async_trload", "qr_async_trload_v3"]:
elif self.pipeline_tag in ["qr", "qs", "qr_async", "qr_async_trload", "qr_async_trload_v3"]:
bk0submax = K0_MAX_SUBMAX_MAP[self.bk0max]
if self.dvpad == "t":
return f"true /*a.hdim_v % {bk0submax} != 0*/" # TODO: order of get_pipelines() matters! (ugly)
@@ -1046,7 +1034,7 @@ class KernelComponentFactoryGfx9(CompatibilityRuleFactoryGfx9):
pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip
else:
pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "f", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr_async", "row", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip
if receipt == 1 and bias != "bias":
pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, skip, "f", sink)) # fmt: skip # TODO: cover arbitraty hdim# fmt: skip
@@ -1060,10 +1048,10 @@ class KernelComponentFactoryGfx9(CompatibilityRuleFactoryGfx9):
["f", "t"],
):
if hdim == 64:
pipelines.append(FmhaFwdPipeline("qr", "row", "t", "f", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr", "row", "f", "f", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip
else:
pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "f", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr_async", "row", "f", "f", "f", "f", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr_async", "row", "t", "t", "t", "t", logits, bias, "f", "f", qscale, mask, "f", "f", sink)) # fmt: skip
return pipelines
@@ -1169,6 +1157,9 @@ class KernelComponentFactoryGfx950(
):
pipelines.append(FmhaFwdPipeline("qr", "col", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr", "col", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip
if hdim > 64 and dtype in cls._DT_MXFP8:
pipelines.append(FmhaFwdPipeline("qr_async", "col", "f", "f", "f", "f", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip
pipelines.append(FmhaFwdPipeline("qr_async", "col", "t", "t", "t", "t", logits, bias, lse, dropout, qscale, mask, "f", "f", sink)) # fmt: skip
return pipelines

2
example/ck_tile/01_fmha/script/benchmark_fwd.sh
View File

@@ -53,4 +53,4 @@ $EXE $base_group_args -s_qpad=1152,896,576,320 -s_kpad=1152,896,576,320
$EXE $base_group_args -s_qpad=1536,1152,768,384 -s_kpad=1536,1152,768,384
# high physical pad
$EXE $base_group_args -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512
$EXE $base_group_args -s_qpad=2048,1536,1024,512 -s_kpad=2048,1536,1024,512

10
include/ck_tile/core/tensor/tile_window.hpp
View File

@@ -506,9 +506,8 @@ struct tile_window_with_static_distribution
using SFC_Ys = typename Traits::SFC_Ys;
static constexpr index_t YElementSize =
typename Base::TileDstr{}.get_ys_to_d_descriptor().get_element_space_size();
static_assert(YElementSize % (Traits::PackedSize * Traits::ScalarPerVector) == 0);
using vectorized_tbuf =
array<vector_t, YElementSize / (Traits::PackedSize * Traits::ScalarPerVector)>;
static_assert(YElementSize % Traits::ScalarPerVector == 0);
using vectorized_tbuf = array<vector_t, YElementSize / Traits::ScalarPerVector>;
constexpr auto tile_dstr = typename Base::TileDstr{};
@@ -534,10 +533,11 @@ struct tile_window_with_static_distribution
constexpr index_t d =
tile_dstr.get_ys_to_d_descriptor().calculate_offset(idx_ys_start) /
Traits::PackedSize;
static_assert(d % Traits::ScalarPerVector == 0);
static_assert(Traits::ScalarPerVector % Traits::PackedSize == 0);
static_assert(d % (Traits::ScalarPerVector / Traits::PackedSize) == 0);
this->get_bottom_tensor_view().template get_vectorized_elements_raw<vector_t>(
dst_vec_tbuf.template at<d / Traits::ScalarPerVector>(),
dst_vec_tbuf.template at<d / (Traits::ScalarPerVector / Traits::PackedSize)>(),
bottom_tensor_thread_coord,
0 /**/,
bool_constant<oob_conditional_check>{},

291
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp
View File

@@ -7,6 +7,7 @@
#include "ck_tile/ops/common/tensor_layout.hpp"
#include "ck_tile/ops/fmha/block/block_attention_bias_enum.hpp"
#include "ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_default_policy.hpp"
#include "ck_tile/ops/fmha/block/cast_tile_mx.hpp"
#include "ck_tile/ops/fmha/block/block_dropout.hpp"
#include "ck_tile/ops/reduce/block/block_reduce.hpp"
@@ -29,6 +30,10 @@ struct BlockFmhaPipelineQRKSVSAsync
using PDataType = remove_cvref_t<typename Problem::PDataType>;
using OaccDataType = remove_cvref_t<typename Problem::OaccDataType>;
using ODataType = remove_cvref_t<typename Problem::ODataType>;
using QScaleDataType = remove_cvref_t<typename Problem::QScaleDataType>;
using KScaleDataType = remove_cvref_t<typename Problem::KScaleDataType>;
using VScaleDataType = remove_cvref_t<typename Problem::VScaleDataType>;
using PScaleDataType = remove_cvref_t<typename Problem::PScaleDataType>;
using AttentionVariant = remove_cvref_t<typename Problem::AttentionVariant>;
using FmhaMask = remove_cvref_t<typename Problem::FmhaMask>;
@@ -50,21 +55,20 @@ struct BlockFmhaPipelineQRKSVSAsync
static_assert(kSubQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");
static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
// TODO: seq_q always support padding, hdim_q/v support multiple of vector(like 8x)
// only need special care about seq_k padding (oob need set -INF of p instead of zero)
static_assert(Problem::kPadSeqLenQ == true && Problem::kPadHeadDimQ == true &&
Problem::kPadHeadDimV == true);
static constexpr bool kPadSeqLenQ = true;
static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ;
static constexpr bool kPadSeqLenK = Problem::kPadSeqLenK;
static constexpr bool kPadHeadDimQ = true; // support multiple of vector(like 8x)
static constexpr bool kPadHeadDimV = true; // support multiple of vector(like 8x)
static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV;
static constexpr bool kHasLogitsSoftCap = Problem::kHasLogitsSoftCap;
static constexpr auto BiasEnum = Problem::BiasEnum;
static constexpr bool kStoreLSE = Problem::kStoreLSE;
static constexpr bool kHasDropout = Problem::kHasDropout;
static constexpr bool kHasSink = Problem::kHasSink;
static constexpr ck_tile::index_t kQKScaleGranularity = Problem::kQKScaleGranularity;
static constexpr ck_tile::index_t kVScaleGranularity = Problem::kVScaleGranularity;
// For BLOCKSCALE: shift value for exp2(x + shift) to scale P to [0, 2^shift]
static constexpr float OCP_FP8_SHIFT = 8.0f;
static constexpr float FNUZ_FP8_SHIFT = 7.0f;
@@ -82,7 +86,8 @@ struct BlockFmhaPipelineQRKSVSAsync
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
return Policy::template GetAlignmentV<Problem>();
else
return kPadSeqLenK ? 1 : Policy::template GetAlignmentV<Problem>();
return kPadSeqLenK ? numeric_traits<VDataType>::PackedSize
: Policy::template GetAlignmentV<Problem>();
}();
static constexpr index_t kAlignmentO = Policy::template GetAlignmentO<Problem>();
static constexpr index_t kAlignmentBias =
@@ -201,9 +206,12 @@ struct BlockFmhaPipelineQRKSVSAsync
const float* k_descale_ptr,
const float* v_descale_ptr,
const index_t block_scale_size_kv,
const QScaleDramBlockWindowTmp&, // M0*(K0/kQKScaleGranularity) tile
const KScaleDramBlockWindowTmp&, // N0*(K0/kQKScaleGranularity) tile
const VScaleDramBlockWindowTmp&, // N1*(K1/kVScaleGranularity) tile
const QScaleDramBlockWindowTmp&
q_scale_dram_block_window_tmp, // M0*(K0/kQKScaleGranularity) tile
const KScaleDramBlockWindowTmp&
k_scale_dram_block_window_tmp, // N0*(K0/kQKScaleGranularity) tile
const VScaleDramBlockWindowTmp&
v_scale_dram_block_window_tmp, // N1*(K1/kVScaleGranularity) tile
const float sink_v) const
{
static_assert(
@@ -213,6 +221,8 @@ struct BlockFmhaPipelineQRKSVSAsync
"wrong!");
static_assert(kM0 == QDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
kSubQKHeaddim ==
QDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] &&
kN0 == KDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
kK0 == KDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] &&
kN1 == VDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
@@ -220,8 +230,28 @@ struct BlockFmhaPipelineQRKSVSAsync
kM0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
kN0 == BiasDramBlockWindowTmp{}.get_window_lengths()[number<1>{}],
"wrong!");
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
static_assert(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::ColumnMajor>);
static_assert(QScaleEnum != BlockAttentionQuantScaleEnum::MX);
static_assert(
std::is_same_v<QScaleDataType,
remove_cvref_t<typename QScaleDramBlockWindowTmp::DataType>> &&
std::is_same_v<KScaleDataType,
remove_cvref_t<typename KScaleDramBlockWindowTmp::DataType>> &&
std::is_same_v<VScaleDataType,
remove_cvref_t<typename VScaleDramBlockWindowTmp::DataType>>);
static_assert(kM0 == QScaleDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
kSubQKHeaddim ==
QScaleDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] *
kQKScaleGranularity &&
kN0 == KScaleDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
kK0 == KScaleDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] *
kQKScaleGranularity &&
kN1 == VScaleDramBlockWindowTmp{}.get_window_lengths()[number<0>{}] &&
kK1 == VScaleDramBlockWindowTmp{}.get_window_lengths()[number<1>{}] *
kVScaleGranularity);
}
constexpr auto LdsSeq = Policy::template GetLdsBufferSequence<Problem>();
@@ -400,12 +430,52 @@ struct BlockFmhaPipelineQRKSVSAsync
{0, kv_load_start}, // TODO: hdim split?
Policy::template MakeVDramTileDistribution<Problem>());
auto q_scale = [&] {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
auto q_scale_dram_window =
make_tile_window(q_scale_dram_block_window_tmp.get_bottom_tensor_view(),
q_scale_dram_block_window_tmp.get_window_lengths(),
q_scale_dram_block_window_tmp.get_window_origin(),
Policy::template MakeQScaleRegTileDistribution<Problem>());
return load_tile(q_scale_dram_window);
}
else
{
return null_tensor{};
}
}();
auto k_scale_dram_block_window = [&] {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
return make_tile_window(k_scale_dram_block_window_tmp.get_bottom_tensor_view(),
k_scale_dram_block_window_tmp.get_window_lengths(),
{seqlen_k_start, 0});
}
else
{
return make_null_tile_window(make_tuple());
}
}();
auto v_scale_dram_window = [&] {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
return make_tile_window(v_scale_dram_block_window_tmp.get_bottom_tensor_view(),
v_scale_dram_block_window_tmp.get_window_lengths(),
{0, seqlen_k_start / kVScaleGranularity},
Policy::template MakeVScaleRegTileDistribution<Problem>());
}
else
{
return make_null_tile_window(make_tuple());
}
}();
// prefetch K tile
async_load_tile_raw(
k_lds_store(LdsSeq.at(number<0>{})), k_dram_window, number<-1>{}, k_oob_ck, k_pre_np);
move_tile_window(k_dram_window, {0, kK0});
__builtin_amdgcn_sched_barrier(0);
buffer_load_fence(k_dram_window.get_num_of_access(), q.get_thread_buffer());
(void)q_element_func; // ??? rocm-6.x if use q element func will have scratch on hdim=64/32
// auto q_tile = q; // tile_elementwise_in(q_element_func, q);
@@ -426,6 +496,56 @@ struct BlockFmhaPipelineQRKSVSAsync
const index_t kv_idx = (kv_load_start + i_total_loops * kN0) / block_scale_size_kv;
k_descale = k_descale_ptr[kv_idx];
}
auto k_scale_dram_window = [&] {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
return make_tile_window(
k_scale_dram_block_window.get_bottom_tensor_view(),
k_scale_dram_block_window.get_window_lengths(),
k_scale_dram_block_window.get_window_origin(),
Policy::template MakeKScaleRegTileDistribution<Problem>());
}
else
{
return make_null_tile_window(make_tuple());
}
}();
auto load_k_scale_block_tile = [&] {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
auto t = load_tile(k_scale_dram_window);
move_tile_window(k_scale_dram_window, {0, kK0 / kQKScaleGranularity});
return t;
}
else
{
return make_null_tile_window(make_tuple());
}
};
auto k_scale_block_tile = load_k_scale_block_tile();
auto run_gemm_0 = [&](auto i_k0) {
auto q_slice =
get_slice_tile(q, sequence<0, i_k0 * kK0>{}, sequence<kM0, (i_k0 + 1) * kK0>{});
auto k_slice =
get_slice_tile(k_lds_load,
sequence<(LdsSeq.at(number<i_k0>{})) * kN0, 0>{},
sequence<(LdsSeq.at(number<i_k0>{}) + 1) * kN0, kK0>{});
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
auto q_scale_slice =
get_slice_tile(q_scale,
sequence<0, i_k0*(kK0 / kQKScaleGranularity)>{},
sequence<kM0, (i_k0 + 1) * (kK0 / kQKScaleGranularity)>{});
gemm_0(s_acc, q_slice, q_scale_slice, k_slice, k_scale_block_tile);
}
else
{
gemm_0(s_acc, q_slice, k_slice);
}
};
// STAGE 1, QK gemm
clear_tile(s_acc); // initialize C
if constexpr(k0_loops > 1)
@@ -442,12 +562,8 @@ struct BlockFmhaPipelineQRKSVSAsync
async_load_fence(k_dram_window.get_num_of_access());
__builtin_amdgcn_s_barrier();
__builtin_amdgcn_sched_barrier(0);
gemm_0(s_acc,
get_slice_tile(
q, sequence<0, i_k0 * kK0>{}, sequence<kM0, (i_k0 + 1) * kK0>{}),
get_slice_tile(k_lds_load,
sequence<(LdsSeq.at(number<i_k0>{})) * kN0, 0>{},
sequence<(LdsSeq.at(number<i_k0>{}) + 1) * kN0, kK0>{}));
run_gemm_0(number<i_k0>{});
k_scale_block_tile = load_k_scale_block_tile();
});
}
@@ -463,13 +579,7 @@ struct BlockFmhaPipelineQRKSVSAsync
auto v_buf = load_tile(v_dram_window, number<-1>{}, bool_constant<false>{});
__builtin_amdgcn_sched_barrier(0);
{ // tail
gemm_0(
s_acc,
get_slice_tile(
q, sequence<0, (k0_loops - 1) * kK0>{}, sequence<kM0, k0_loops * kK0>{}),
get_slice_tile(k_lds_load,
sequence<(LdsSeq.at(number<k0_loops - 1>{})) * kN0, 0>{},
sequence<(LdsSeq.at(number<k0_loops - 1>{}) + 1) * kN0, kK0>{}));
run_gemm_0(number<k0_loops - 1>{});
}
__builtin_amdgcn_sched_barrier(1);
// dequant
@@ -638,11 +748,11 @@ struct BlockFmhaPipelineQRKSVSAsync
tile_elementwise_in(v_element_func, v_buf)); // store the prefetch
}
move_tile_window(
v_dram_window,
{0, kK1}); // will have scratch if move this right after load_tile(v_dram)...
if constexpr(k1_loops > 1)
{
move_tile_window(
v_dram_window,
{0, kK1}); // will have scratch if move this right after load_tile(v_dram)...
v_buf = load_tile(
v_dram_window, number<-1>{}, bool_constant<false>{}); // load next v_buf
}
@@ -828,21 +938,19 @@ struct BlockFmhaPipelineQRKSVSAsync
randval_ptr, seq_offset, p_compute, randval_dram_window);
}
const auto p = [&]() {
#if CK_TILE_FMHA_FLOAT_TO_FLOAT16_RTN
// For fp32 to fp16,
// impl::cast_tile_pkrtz_fp16_fp32 would cause precision issue,
// since it uses __builtin_amdgcn_cvt_pkrtz, which is round to zero.
return cast_tile<PDataType>(tile_elementwise_in(p_compute_element_func, p_compute));
#else
if constexpr(std::is_same_v<PDataType, fp16_t>)
return impl::cast_tile_pkrtz_fp16_fp32<PDataType>(
tile_elementwise_in(p_compute_element_func, p_compute));
auto load_v_scale_block_tile = [&] {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
auto t = load_tile(v_scale_dram_window);
move_tile_window(v_scale_dram_window, {0, kK1 / kVScaleGranularity});
return t;
}
else
return cast_tile<PDataType>(
tile_elementwise_in(p_compute_element_func, p_compute));
#endif
}();
{
return make_null_tile_window(make_tuple());
}
};
auto v_scale_block_tile = load_v_scale_block_tile();
float v_descale = 1.0f;
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE)
@@ -851,6 +959,48 @@ struct BlockFmhaPipelineQRKSVSAsync
const index_t kv_idx = (kv_load_start + i_total_loops * kN0) / block_scale_size_kv;
v_descale = v_descale_ptr[kv_idx];
}
const auto p_p_scale = [&] {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
auto p_result = make_static_distributed_tensor<PDataType>(
p_compute.get_tile_distribution());
auto p_scale_result = make_static_distributed_tensor<PScaleDataType>(
Policy::template MakePScaleRegTileDistribution<Problem>());
constexpr auto config =
decltype(gemm_1)::Policy::template GetWarpGemmMWarpNWarp<Problem>();
using WG = remove_cvref_t<decltype(config.template at<0>())>;
cast_tile_mx<kVScaleGranularity, WG::WarpGemmAttribute::Impl::kAMLane>(
p_result, p_scale_result, p_compute);
return make_tuple(p_result, p_scale_result);
}
else
{
const auto p_result = [&]() {
#if CK_TILE_FMHA_FLOAT_TO_FLOAT16_RTN
// For fp32 to fp16,
// impl::cast_tile_pkrtz_fp16_fp32 would cause precision issue,
// since it uses __builtin_amdgcn_cvt_pkrtz, which is round to zero.
return cast_tile<PDataType>(
tile_elementwise_in(p_compute_element_func, p_compute));
#else
if constexpr(std::is_same_v<PDataType, fp16_t>)
return impl::cast_tile_pkrtz_fp16_fp32<PDataType>(
tile_elementwise_in(p_compute_element_func, p_compute));
else
return cast_tile<PDataType>(
tile_elementwise_in(p_compute_element_func, p_compute));
#endif
}();
return make_tuple(p_result, null_tensor{});
}
}();
const auto p = p_p_scale[number<0>{}];
const auto p_scale = p_p_scale[number<1>{}];
// STAGE 3, KV gemm
auto o_acc0 = decltype(o_acc){};
clear_tile(o_acc0);
@@ -865,22 +1015,42 @@ struct BlockFmhaPipelineQRKSVSAsync
return o_acc;
}
}();
auto run_gemm_1 = [&](auto i_k1) {
auto p_slice =
get_slice_tile(p, sequence<0, i_k1 * kK1>{}, sequence<kM0, (i_k1 + 1) * kK1>{});
auto v_slice = get_slice_tile(
v_lds_window,
sequence<(LdsSeq.at(number<k0_loops + i_k1>{})) * kN1, 0>{},
sequence<(LdsSeq.at(number<k0_loops + i_k1>{}) + 1) * kN1, kK1>{});
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
auto p_scale_slice =
get_slice_tile(p_scale,
sequence<0, i_k1*(kK1 / kVScaleGranularity)>{},
sequence<kM0, (i_k1 + 1) * (kK1 / kVScaleGranularity)>{});
gemm_1(o_acc_, p_slice, p_scale_slice, v_slice, v_scale_block_tile);
}
else if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE)
{
gemm_1(o_acc0, p_slice, v_slice);
}
else
{
gemm_1(o_acc_, p_slice, v_slice);
}
};
if constexpr(k1_loops > 1)
{
static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) {
if constexpr(i_k1 != 0 && i_k1 < k1_loops - 1)
if constexpr(i_k1 != 0)
{
v_buf = load_tile(
v_dram_window, number<-1>{}, bool_constant<false>{}); // load next v_buf
}
block_sync_lds();
gemm_1(o_acc_,
get_slice_tile(
p, sequence<0, i_k1 * kK1>{}, sequence<kM0, (i_k1 + 1) * kK1>{}),
get_slice_tile(
v_lds_window,
sequence<(LdsSeq.at(number<k0_loops + i_k1>{})) * kN1, 0>{},
sequence<(LdsSeq.at(number<k0_loops + i_k1>{}) + 1) * kN1, kK1>{}));
run_gemm_1(number<i_k1>{});
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
{
@@ -904,8 +1074,8 @@ struct BlockFmhaPipelineQRKSVSAsync
store_tile(v_lds_window_tmp,
tile_elementwise_in(v_element_func, v_buf)); // store next v_buf
}
if constexpr(i_k1 < k1_loops - 1)
move_tile_window(v_dram_window, {0, kK1});
move_tile_window(v_dram_window, {0, kK1});
v_scale_block_tile = load_v_scale_block_tile();
});
}
i_total_loops++;
@@ -913,6 +1083,7 @@ struct BlockFmhaPipelineQRKSVSAsync
{
if constexpr(kHasSink)
{
// TODO: this never happens because of i_total_loops++
if(i_total_loops == 0)
{
move_tile_window(k_dram_block_window, {seqlen_k_start - sink_seq_end, 0});
@@ -920,6 +1091,10 @@ struct BlockFmhaPipelineQRKSVSAsync
}
}
move_tile_window(k_dram_block_window, {kN0, 0});
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::MX)
{
move_tile_window(k_scale_dram_block_window, {kN0, 0});
}
k_dram_window.set_window_origin(k_dram_block_window.get_window_origin());
@@ -936,13 +1111,7 @@ struct BlockFmhaPipelineQRKSVSAsync
// tail
{
block_sync_lds();
gemm_1(
o_acc_,
get_slice_tile(p, sequence<0, (k1_loops - 1) * kK1>{}, sequence<kM0, kN0>{}),
get_slice_tile(
v_lds_window,
sequence<(LdsSeq.at(number<k0_loops + k1_loops - 1>{})) * kN1, 0>{},
sequence<(LdsSeq.at(number<k0_loops + k1_loops - 1>{}) + 1) * kN1, kK1>{}));
run_gemm_1(number<k1_loops - 1>{});
}
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE)

86
test/ck_tile/fmha/test_fmha_fwd.cpp
View File

@@ -49,6 +49,10 @@ struct TestConfigs
static constexpr bool def_is_v_rowmajor = true;
static constexpr auto init_method = "uf";
static int adjust_seqlen(int seqlen) { return seqlen; }
static int adjust_hdim(int hdim)
{
return hdim < 0 ? hdim : ck_tile::integer_least_multiple(hdim, 8);
}
};
template <>
@@ -67,6 +71,10 @@ struct TestConfigs<FmhaFwdFp8Bf16>
// When there are no fp8 instances with padding, pad seqlen to avoid skipping most of the tests:
// return ck_tile::integer_least_multiple(seqlen, 128);
static int adjust_seqlen(int seqlen) { return seqlen; }
static int adjust_hdim(int hdim)
{
return hdim < 0 ? hdim : ck_tile::integer_least_multiple(hdim, 16);
}
};
template <>
@@ -82,6 +90,10 @@ struct TestConfigs<FmhaFwdMxFp8>
static constexpr bool def_is_v_rowmajor = false;
static constexpr auto init_method = "3";
static int adjust_seqlen(int seqlen) { return seqlen; }
static int adjust_hdim(int hdim)
{
return hdim < 0 ? hdim : ck_tile::integer_least_multiple(hdim, 16);
}
};
template <>
@@ -100,6 +112,10 @@ struct TestConfigs<FmhaFwdMxFp4>
{
return seqlen < 0 ? seqlen : ck_tile::integer_least_multiple(seqlen, 2);
}
static int adjust_hdim(int hdim)
{
return hdim < 0 ? hdim : ck_tile::integer_least_multiple(hdim, 32);
}
};
template <>
@@ -123,6 +139,7 @@ struct TestConfigs<FmhaFwdFp32>
static constexpr bool def_is_v_rowmajor = true;
static constexpr auto init_method = "uf";
static int adjust_seqlen(int seqlen) { return seqlen; }
static int adjust_hdim(int hdim) { return hdim; }
};
static auto HDimValues = ValuesIn(TestConfigs<DataTypeConfig>::HDimValues);
@@ -135,6 +152,7 @@ constexpr bool def_lse = TestConfigs<DataTypeConfig>::def_lse;
constexpr bool def_is_v_rowmajor = TestConfigs<DataTypeConfig>::def_is_v_rowmajor;
constexpr auto init_method = TestConfigs<DataTypeConfig>::init_method;
int adjust_seqlen(int seqlen) { return TestConfigs<DataTypeConfig>::adjust_seqlen(seqlen); }
int adjust_hdim(int hdim) { return TestConfigs<DataTypeConfig>::adjust_hdim(hdim); }
// Random seed used for initializing input tensors. 0 for non-deterministic seed
CK_TILE_DECLARE_ENV_VAR(CK_TILE_TEST_SEED, uint64_t, 123456)
@@ -253,6 +271,70 @@ TEST_P(AllLong, DataTypeConfig)
CHECK_RESULT(result);
}
class General
: public TestWithParam<std::tuple<std::tuple<int, int>,
bool,
bool,
mode_enum,
std::tuple<int, int, int, int, int, int, std::string>>>
{
};
INSTANTIATE_TEST_SUITE_P(TestCkTileFmhaFwd,
General,
Combine(HDimValues,
Bool(),
IsVRowmajorValues,
ModeValues,
Values(std::tuple{2, 2, 1, 55, 256, -1, "0"},
std::tuple{1, 3, -1, 100, 51, -1, "0"},
std::tuple{2, 1, -1, 99, 256, -1, "1"},
std::tuple{1, 2, 1, 1024, 256, -1, "2"},
std::tuple{2, 1, -1, 3, 99, -1, "2"},
std::tuple{1, 2, 1, 33, 0, -1, "2"},
std::tuple{1, 2, 1, 1, 10, 32, "2"})));
TEST_P(General, DataTypeConfig)
{
auto [hdims, perm, is_v_rowmajor, mode, dims_mask] = GetParam();
auto [hdim_q, hdim_v] = hdims;
auto [batch, nhead, nhead_k, seqlen_q, seqlen_k, seqlen_kpad, mask_str] = dims_mask;
auto result = fmha_fwd_run<DataTypeConfig>(mode,
batch,
nhead,
nhead_k,
{adjust_seqlen(seqlen_q)},
{adjust_seqlen(seqlen_k)},
hdim_q,
hdim_v,
0, // seqlen_knew
{-1}, // seqlen_qpads
{seqlen_kpad}, // seqlen_kpads
{}, // q_eff_lens_per_batch
{}, // kv_eff_lens_per_batch
0, // rotary_dim
perm, // i_perm
perm, // o_perm
0, // scale_s
0, // logits_soft_cap
is_v_rowmajor, // is_v_rowmajor
def_lse, // lse
0, // page_block_size
false, // use_cache_batch_idx
"n", // bias_str
0.0f, // p_drop
0, // drop_seed
0, // drop_offset
false, // drop_prefs
mask_str, // mask_str
qscale_str,
true, // is_rotary_interleaved
1, // num_splits
COMMON_ARGS);
CHECK_RESULT(result);
}
// ---------------------------------------------------------------
// Negative tests: padding not supported with appendkv/splitkv/pagedkv
// ---------------------------------------------------------------
@@ -430,8 +512,8 @@ TEST_P(HDimPadding, DataTypeConfig)
nhead_k,
{adjust_seqlen(seqlen_q)},
{adjust_seqlen(seqlen_k)},
hdim_q,
hdim_v,
adjust_hdim(hdim_q),
adjust_hdim(hdim_v),
0, // seqlen_knew
{-1}, // seqlen_qpads
{seqlen_kpad}, // seqlen_kpads