ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 09:38:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

[rocm-libraries] ROCm/rocm-libraries#9461 (commit 0f33948)

Browse Source 
Revert PR #6526 "feat: [CK Tile] mxfp8 support for qr async
 pipeline (#6526)" (#9461)

This reverts commit 3e0171014f5d0e31152376271b6208c6587ef786.

## Motivation

The PR is identified breaking the CI RUN_CK_TILE_FMHA_TESTS.

http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2fComposable%20Kernel/detail/develop/1762/pipeline/305

and Aiter test_mha.py
http://micimaster.amd.com/blue/organizations/jenkins/rocm-libraries-folder%2FComposable%20Kernel/detail/users%2Fjialuo%2Fck%2Ffix_ci_aiter_error/33/pipeline/246

## Technical Details

<!-- Explain the changes along with any relevant GitHub links. -->

## Test Plan

<!-- Explain any relevant testing done to verify this PR. -->

## Test Result

<!-- Briefly summarize test outcomes. -->

## Submission Checklist

- [ ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.
This commit is contained in:
JiaLuo-CAN
2026-07-16 15:35:37 +00:00
committed by assistant-librarian[bot]
parent 6648115aed
commit 9d8386b920
5 changed files with 88 additions and 330 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -506,8 +506,9 @@ 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::ScalarPerVector == 0);
using vectorized_tbuf = array<vector_t, YElementSize / Traits::ScalarPerVector>;
static_assert(YElementSize % (Traits::PackedSize * Traits::ScalarPerVector) == 0);
using vectorized_tbuf =
array<vector_t, YElementSize / (Traits::PackedSize * Traits::ScalarPerVector)>;
constexpr auto tile_dstr = typename Base::TileDstr{};
@@ -533,11 +534,10 @@ 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(Traits::ScalarPerVector % Traits::PackedSize == 0);
static_assert(d % (Traits::ScalarPerVector / Traits::PackedSize) == 0);
static_assert(d % Traits::ScalarPerVector == 0);
this->get_bottom_tensor_view().template get_vectorized_elements_raw<vector_t>(
dst_vec_tbuf.template at<d / (Traits::ScalarPerVector / Traits::PackedSize)>(),
dst_vec_tbuf.template at<d / Traits::ScalarPerVector>(),
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,7 +7,6 @@
#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"
@@ -30,10 +29,6 @@ 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>;
@@ -55,20 +50,21 @@ struct BlockFmhaPipelineQRKSVSAsync
static_assert(kSubQKHeaddim <= 256, "hdim bigger than 256 is not suitable for this pipeline!");
static constexpr bool kIsGroupMode = Problem::kIsGroupMode;
static constexpr bool kPadSeqLenQ = Problem::kPadSeqLenQ;
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 kPadSeqLenK = Problem::kPadSeqLenK;
static constexpr bool kPadHeadDimQ = Problem::kPadHeadDimQ;
static constexpr bool kPadHeadDimV = Problem::kPadHeadDimV;
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 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;
@@ -86,8 +82,7 @@ struct BlockFmhaPipelineQRKSVSAsync
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
return Policy::template GetAlignmentV<Problem>();
else
return kPadSeqLenK ? numeric_traits<VDataType>::PackedSize
: Policy::template GetAlignmentV<Problem>();
return kPadSeqLenK ? 1 : Policy::template GetAlignmentV<Problem>();
}();
static constexpr index_t kAlignmentO = Policy::template GetAlignmentO<Problem>();
static constexpr index_t kAlignmentBias =
@@ -206,12 +201,9 @@ struct BlockFmhaPipelineQRKSVSAsync
const float* k_descale_ptr,
const float* v_descale_ptr,
const index_t block_scale_size_kv,
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 QScaleDramBlockWindowTmp&, // M0*(K0/kQKScaleGranularity) tile
const KScaleDramBlockWindowTmp&, // N0*(K0/kQKScaleGranularity) tile
const VScaleDramBlockWindowTmp&, // N1*(K1/kVScaleGranularity) tile
const float sink_v) const
{
static_assert(
@@ -221,8 +213,6 @@ 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>{}] &&
@@ -230,28 +220,8 @@ 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(
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);
}
static_assert(QScaleEnum != BlockAttentionQuantScaleEnum::MX);
constexpr auto LdsSeq = Policy::template GetLdsBufferSequence<Problem>();
@@ -430,52 +400,12 @@ 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);
@@ -496,56 +426,6 @@ 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)
@@ -562,8 +442,12 @@ struct BlockFmhaPipelineQRKSVSAsync
async_load_fence(k_dram_window.get_num_of_access());
__builtin_amdgcn_s_barrier();
__builtin_amdgcn_sched_barrier(0);
run_gemm_0(number<i_k0>{});
k_scale_block_tile = load_k_scale_block_tile();
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>{}));
});
}
@@ -579,7 +463,13 @@ struct BlockFmhaPipelineQRKSVSAsync
auto v_buf = load_tile(v_dram_window, number<-1>{}, bool_constant<false>{});
__builtin_amdgcn_sched_barrier(0);
{ // tail
run_gemm_0(number<k0_loops - 1>{});
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>{}));
}
__builtin_amdgcn_sched_barrier(1);
// dequant
@@ -748,11 +638,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
}
@@ -938,19 +828,21 @@ struct BlockFmhaPipelineQRKSVSAsync
randval_ptr, seq_offset, p_compute, randval_dram_window);
}
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;
}
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));
else
{
return make_null_tile_window(make_tuple());
}
};
auto v_scale_block_tile = load_v_scale_block_tile();
return cast_tile<PDataType>(
tile_elementwise_in(p_compute_element_func, p_compute));
#endif
}();
float v_descale = 1.0f;
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE)
@@ -959,48 +851,6 @@ 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);
@@ -1015,42 +865,22 @@ 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)
if constexpr(i_k1 != 0 && i_k1 < k1_loops - 1)
{
v_buf = load_tile(
v_dram_window, number<-1>{}, bool_constant<false>{}); // load next v_buf
}
block_sync_lds();
run_gemm_1(number<i_k1>{});
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>{}));
if constexpr(std::is_same_v<VLayout, ck_tile::tensor_layout::gemm::RowMajor>)
{
@@ -1074,8 +904,8 @@ struct BlockFmhaPipelineQRKSVSAsync
store_tile(v_lds_window_tmp,
tile_elementwise_in(v_element_func, v_buf)); // store next v_buf
}
move_tile_window(v_dram_window, {0, kK1});
v_scale_block_tile = load_v_scale_block_tile();
if constexpr(i_k1 < k1_loops - 1)
move_tile_window(v_dram_window, {0, kK1});
});
}
i_total_loops++;
@@ -1083,7 +913,6 @@ 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});
@@ -1091,10 +920,6 @@ 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());
@@ -1111,7 +936,13 @@ struct BlockFmhaPipelineQRKSVSAsync
// tail
{
block_sync_lds();
run_gemm_1(number<k1_loops - 1>{});
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>{}));
}
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::BLOCKSCALE)