ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 10:08:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

kernel and pipeline add q per-tensor kv block scale

Browse Source
This commit is contained in:
ltqin
2026-01-29 07:16:47 +00:00
parent 9b168082b7
commit 9a657edaaf
3 changed files with 253 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
include/ck_tile/ops/fmha/block/block_attention_quant_scale_enum.hpp
View File

@@ -13,6 +13,7 @@ enum class BlockAttentionQuantScaleEnum
NO_SCALE = 0,
PERTENSOR = 1,
BLOCKSCALE,
Q_PERTENSOR_KV_BLOCKSCALE,
};
template <BlockAttentionQuantScaleEnum>
@@ -33,5 +34,10 @@ struct BlockAttentionQuantScaleEnumToStr<BlockAttentionQuantScaleEnum::BLOCKSCAL
{
static constexpr const char* name = "blockscale";
};
template <>
struct BlockAttentionQuantScaleEnumToStr<BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE>
{
static constexpr const char* name = "q_pertensor_kv_blockscale";
};
} // namespace ck_tile

167
include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp
View File

@@ -192,6 +192,25 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
const void* v_descale_ptr = nullptr;
};
struct FmhaFwdCommonQPerTensorKVBlockScaleKargs : public FmhaFwdCommonQScaleKargs
{
ck_tile::index_t nhead_stride_k_descale;
ck_tile::index_t nhead_stride_v_descale;
ck_tile::index_t block_scale_size_kv;
};
struct FmhaFwdBatchQPerTensorKVBlockScaleKargs : public FmhaFwdCommonQPerTensorKVBlockScaleKargs
{
ck_tile::index_t batch_stride_k_descale;
ck_tile::index_t batch_stride_v_descale;
};
struct FmhaFwdGroupQPerTensorKVBlockScaleKargs : public FmhaFwdCommonQPerTensorKVBlockScaleKargs
{
const int32_t* block_scale_seqstart_k_ptr;
};
struct FmhaFwdDropoutSeedOffset
{
template <typename T>
@@ -255,9 +274,13 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
FmhaFwdEmptyKargs<0>>>,
std::conditional_t<kHasMask, FmhaFwdMaskKargs, FmhaFwdEmptyKargs<1>>,
std::conditional_t<kStoreLSE, FmhaFwdCommonLSEKargs, FmhaFwdEmptyKargs<2>>,
std::conditional_t<QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR,
FmhaFwdCommonQScaleKargs,
FmhaFwdEmptyKargs<3>>,
std::conditional_t<
QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR,
FmhaFwdCommonQScaleKargs,
std::conditional_t<QScaleEnum ==
BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE,
FmhaFwdBatchQPerTensorKVBlockScaleKargs,
FmhaFwdEmptyKargs<3>>>,
std::conditional_t<kHasDropout, FmhaFwdBatchModeDropoutKargs, FmhaFwdEmptyKargs<4>>,
std::conditional_t<kHasLogitsSoftCap, FmhaFwdLogitsSoftCapKargs, FmhaFwdEmptyKargs<5>>
{
@@ -276,9 +299,13 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
FmhaFwdEmptyKargs<0>>>,
std::conditional_t<kHasMask, FmhaFwdMaskKargs, FmhaFwdEmptyKargs<1>>,
std::conditional_t<kStoreLSE, FmhaFwdCommonLSEKargs, FmhaFwdEmptyKargs<2>>,
std::conditional_t<QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR,
FmhaFwdCommonQScaleKargs,
FmhaFwdEmptyKargs<3>>,
std::conditional_t<
QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR,
FmhaFwdCommonQScaleKargs,
std::conditional_t<QScaleEnum ==
BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE,
FmhaFwdGroupQPerTensorKVBlockScaleKargs,
FmhaFwdEmptyKargs<3>>>,
std::conditional_t<kHasDropout, FmhaFwdCommonDropoutKargs, FmhaFwdEmptyKargs<4>>,
std::conditional_t<kHasLogitsSoftCap, FmhaFwdLogitsSoftCapKargs, FmhaFwdEmptyKargs<5>>
{
@@ -348,7 +375,12 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
bool s_randval,
std::variant<std::pair<uint64_t, uint64_t>, std::pair<const void*, const void*>>
drop_seed_offset,
const void* sink_ptr = nullptr)
const void* sink_ptr = nullptr,
ck_tile::index_t nhead_stride_k_descale = 0,
ck_tile::index_t nhead_stride_v_descale = 0,
ck_tile::index_t batch_stride_k_descale = 0,
ck_tile::index_t batch_stride_v_descale = 0,
ck_tile::index_t block_scale_size_kv = 0)
{
Kargs kargs{{q_ptr,
k_ptr,
@@ -419,6 +451,20 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
kargs.k_descale_ptr = k_descale_ptr;
kargs.v_descale_ptr = v_descale_ptr;
}
else if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
kargs.q_descale_ptr = q_descale_ptr;
kargs.k_descale_ptr = k_descale_ptr;
kargs.v_descale_ptr = v_descale_ptr;
kargs.nhead_stride_k_descale = nhead_stride_k_descale;
kargs.nhead_stride_v_descale = nhead_stride_v_descale;
kargs.batch_stride_k_descale = batch_stride_k_descale;
kargs.batch_stride_v_descale = batch_stride_v_descale;
kargs.block_scale_size_kv = block_scale_size_kv;
}
if constexpr(kHasDropout)
{
if(drop_seed_offset.index() == 0) // seed & offset come from host
@@ -495,7 +541,13 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
bool s_randval,
std::variant<std::pair<uint64_t, uint64_t>, std::pair<const void*, const void*>>
drop_seed_offset,
const void* sink_ptr = nullptr)
const void* sink_ptr = nullptr,
ck_tile::index_t nhead_stride_k_descale = 0,
ck_tile::index_t nhead_stride_v_descale = 0,
ck_tile::index_t batch_stride_k_descale = 0,
ck_tile::index_t batch_stride_v_descale = 0,
ck_tile::index_t block_scale_size_kv = 0,
const int32_t* block_scale_seqstart_k_ptr = nullptr)
{
Kargs kargs{{q_ptr,
k_ptr,
@@ -563,6 +615,21 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
kargs.k_descale_ptr = k_descale_ptr;
kargs.v_descale_ptr = v_descale_ptr;
}
else if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
kargs.q_descale_ptr = q_descale_ptr;
kargs.k_descale_ptr = k_descale_ptr;
kargs.v_descale_ptr = v_descale_ptr;
kargs.nhead_stride_k_descale = nhead_stride_k_descale;
kargs.nhead_stride_v_descale = nhead_stride_v_descale;
kargs.batch_stride_k_descale = batch_stride_k_descale;
kargs.batch_stride_v_descale = batch_stride_v_descale;
kargs.block_scale_size_kv = block_scale_size_kv;
kargs.block_scale_seqstart_k_ptr = block_scale_seqstart_k_ptr;
}
if constexpr(kHasDropout)
{
if(drop_seed_offset.index() == 0) // seed & offset come from host
@@ -819,6 +886,46 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
static_cast<long_index_t>(i_nhead) * kargs.nhead_stride_o +
batch_offset_o;
// Prepare descale pointers (only used for Q_PERTENSOR_KV_BLOCKSCALE)
[[maybe_unused]] const float* k_descale_ptr = nullptr;
[[maybe_unused]] const float* v_descale_ptr = nullptr;
[[maybe_unused]] index_t block_scale_size_kv = 0;
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
if constexpr(kIsGroupMode)
{
// For group mode, use seqstart_k_ptr to calculate offset
k_descale_ptr =
reinterpret_cast<const float*>(kargs.k_descale_ptr) +
static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) *
kargs.nhead_stride_k_descale +
kargs.batch_stride_k_descale * (kargs.block_scale_seqstart_k_ptr
? kargs.block_scale_seqstart_k_ptr[i_batch]
: 0);
v_descale_ptr =
reinterpret_cast<const float*>(kargs.v_descale_ptr) +
static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) *
kargs.nhead_stride_v_descale +
kargs.batch_stride_v_descale * (kargs.block_scale_seqstart_k_ptr
? kargs.block_scale_seqstart_k_ptr[i_batch]
: 0);
}
else
{
// For batch mode
k_descale_ptr = reinterpret_cast<const float*>(kargs.k_descale_ptr) +
static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) *
kargs.nhead_stride_k_descale +
static_cast<long_index_t>(i_batch) * kargs.batch_stride_k_descale;
v_descale_ptr = reinterpret_cast<const float*>(kargs.v_descale_ptr) +
static_cast<long_index_t>(i_nhead / kargs.nhead_ratio_qk) *
kargs.nhead_stride_v_descale +
static_cast<long_index_t>(i_batch) * kargs.batch_stride_v_descale;
}
block_scale_size_kv = kargs.block_scale_size_kv;
}
// Q/K/V DRAM and DRAM window
const auto q_dram = [&]() {
const auto q_dram_naive = make_naive_tensor_view<address_space_enum::global>(
@@ -1162,6 +1269,13 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
return kargs.scale_s * q_descale * k_descale;
}
else if constexpr(QScaleEnum ==
BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
// Q uses per-tensor scale, K uses block scale (handled in pipeline)
float q_descale = *(reinterpret_cast<const float*>(kargs.q_descale_ptr));
return kargs.scale_s * q_descale;
}
else
{
return kargs.scale_s;
@@ -1235,10 +1349,15 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
kargs.batch_stride_k,
kargs.batch_stride_v,
dropout,
nullptr, // k_descale_ptr not used for PERTENSOR
nullptr, // v_descale_ptr not used for PERTENSOR (handled via
// o_acc_element_func)
0, // block_scale_size_kv not used for PERTENSOR
sink_value);
}
else
else if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
// Q: per-tensor (handled in scale_s), K/V: block scale (handled in pipeline)
return FmhaPipeline{}(q_dram_window,
k_dram_window,
v_dram_window,
@@ -1258,6 +1377,36 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel
kargs.batch_stride_k,
kargs.batch_stride_v,
dropout,
k_descale_ptr,
v_descale_ptr,
block_scale_size_kv,
sink_value);
}
else
{
// NO_SCALE
return FmhaPipeline{}(q_dram_window,
k_dram_window,
v_dram_window,
bias_dram_window,
randval_dram_window,
lse_dram_window,
mask,
position_encoding,
variant_params.sm_scale,
variant,
variant_params,
block_indices,
smem_ptr,
page_idx,
stride_k_for_pipeline,
stride_v_for_pipeline,
kargs.batch_stride_k,
kargs.batch_stride_v,
dropout,
nullptr, // k_descale_ptr not used for NO_SCALE
nullptr, // v_descale_ptr not used for NO_SCALE
0, // block_scale_size_kv not used for NO_SCALE
sink_value);
}
}();

100
include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp
View File

@@ -341,6 +341,9 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
const index_t page_stride_k,
const index_t page_stride_v,
DropoutType& dropout,
const float* k_descale_ptr,
const float* v_descale_ptr,
const index_t block_scale_size_kv,
const float sink_v) const
{
static_assert(
@@ -717,6 +720,13 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
// main loop
do
{
float k_descale = 1.0f;
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
// K and V share the same seqlen_k position within a block
const index_t kv_idx = (seqlen_k_start + i_total_loops * kN0) / block_scale_size_kv;
k_descale = k_descale_ptr[kv_idx];
}
// STAGE 1, QK gemm
clear_tile(s_acc); // initialize C
if constexpr(k0_loops > 1)
@@ -761,6 +771,15 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
sequence<(LdsSeq.at(number<k0_loops - 1>{}) + 1) * kN0, kK0>{}));
}
__builtin_amdgcn_sched_barrier(1);
// dequant
auto s_acc_element_func_ = [&s_acc_element_func, k_descale]() {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
return s_acc_element_func * k_descale;
}
else
return s_acc_element_func;
}();
auto v_buf = load_tile(v_dram_window, number<-1>{}, bool_constant<false>{});
update_v_offsets(number<kK1>{});
@@ -772,7 +791,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
// STAGE 2, scale_s, add bias, mask, softmax
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS)
{
s_acc = tile_elementwise_in(s_acc_element_func, s_acc);
s_acc = tile_elementwise_in(s_acc_element_func_, s_acc);
tile_elementwise_inout([&scale_s](auto& x) { x = x * scale_s; }, s_acc);
tile_elementwise_inout(
[&](auto& x, const auto& y) {
@@ -790,7 +809,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
{
const auto k_origin = k_dram_block_window.get_window_origin();
constexpr auto s_spans = decltype(s_acc)::get_distributed_spans();
s_acc = tile_elementwise_in(s_acc_element_func, s_acc);
s_acc = tile_elementwise_in(s_acc_element_func_, s_acc);
sweep_tile_span(s_spans[number<0>{}], [&](auto idx0) {
sweep_tile_span(s_spans[number<1>{}], [&](auto idx1) {
const auto tile_idx = get_x_indices_from_distributed_indices(
@@ -807,7 +826,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
}
else
{
s_acc = tile_elementwise_in(s_acc_element_func, s_acc);
s_acc = tile_elementwise_in(s_acc_element_func_, s_acc);
if constexpr(kHasLogitsSoftCap)
{
auto apply_logits_transform = [&variant, &variant_params, &block_indices](
@@ -953,7 +972,23 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) {
constexpr auto i_idx = make_tuple(idx0);
#if CK_TILE_FMHA_FWD_FAST_EXP2
auto row_max = scale_s * get_validated_m(m[i_idx]);
// For Q_PERTENSOR_KV_BLOCKSCALE: precompute (m - shift) once per row
// Bias/Alibi/SoftCap: exp2(s - m + shift) = exp2(s - (m - shift))
// else: exp2(scale_s*s - scale_s*m + shift) = exp2(scale_s*s - (scale_s*m -
// shift))
auto validated_m = get_validated_m(m[i_idx]);
auto row_max = scale_s * validated_m;
if constexpr(QScaleEnum ==
BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
#if CK_TILE_USE_OCP_FP8
validated_m -= OCP_FP8_SHIFT; // for Bias/Alibi/SoftCap
row_max -= OCP_FP8_SHIFT; // for else branch
#else
validated_m -= FNUZ_FP8_SHIFT;
row_max -= FNUZ_FP8_SHIFT;
#endif
}
#endif
sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) {
constexpr auto i_j_idx = make_tuple(idx0, idx1);
@@ -961,13 +996,13 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
BiasEnum == BlockAttentionBiasEnum::ALIBI)
{
p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m[i_idx]));
p_compute(i_j_idx) = exp2(s[i_j_idx] - validated_m);
}
else
{
if constexpr(kHasLogitsSoftCap)
{
p_compute(i_j_idx) = exp2(s[i_j_idx] - get_validated_m(m[i_idx]));
p_compute(i_j_idx) = exp2(s[i_j_idx] - validated_m);
}
else
{
@@ -990,20 +1025,34 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
constexpr auto i_idx = make_tuple(idx0);
#if CK_TILE_FMHA_FWD_FAST_EXP2
const auto tmp = [&]() {
auto validated_m_current = get_validated_m(m[i_idx]);
auto row_max =
scale_s * validated_m_current; // Calculate BEFORE applying shift
if constexpr(QScaleEnum ==
BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
#if CK_TILE_USE_OCP_FP8
validated_m_current -= OCP_FP8_SHIFT; // for Bias/Alibi/SoftCap
row_max -= OCP_FP8_SHIFT; // for else branch
#else
validated_m_current -= FNUZ_FP8_SHIFT;
row_max -= FNUZ_FP8_SHIFT;
#endif
}
if constexpr(BiasEnum == BlockAttentionBiasEnum::ELEMENTWISE_BIAS ||
BiasEnum == BlockAttentionBiasEnum::ALIBI)
{
return exp2(m_old[i_idx] - get_validated_m(m[i_idx]));
return exp2(m_old[i_idx] - validated_m_current);
}
else
{
if constexpr(kHasLogitsSoftCap)
{
return exp2(m_old[i_idx] - get_validated_m(m[i_idx]));
return exp2(m_old[i_idx] - validated_m_current);
}
else
{
auto row_max = scale_s * get_validated_m(m[i_idx]);
return exp2(scale_s * m_old[i_idx] - row_max);
}
}
@@ -1048,7 +1097,27 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
#endif
}();
float v_descale = 1.0f;
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
// K and V share the same seqlen_k position within a block
const index_t kv_idx = (seqlen_k_start + i_total_loops * kN0) / block_scale_size_kv;
v_descale = v_descale_ptr[kv_idx];
}
// STAGE 3, KV gemm
auto o_acc0 = decltype(o_acc){};
clear_tile(o_acc0);
auto& o_acc_ = [&o_acc0, &o_acc]() -> auto& {
if constexpr(QScaleEnum == BlockAttentionQuantScaleEnum::Q_PERTENSOR_KV_BLOCKSCALE)
{
return o_acc0;
}
else
{
return o_acc;
}
}();
if constexpr(k1_loops > 1)
{
static_for<0, k1_loops - 1, 1>{}([&](auto i_k1) {
@@ -1060,7 +1129,7 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
v_dram_window.update_page_idx(v_offsets);
}
block_sync_lds();
gemm_1(o_acc,
gemm_1(o_acc_,
get_slice_tile(
p, sequence<0, i_k1 * kK1>{}, sequence<kM0, (i_k1 + 1) * kK1>{}),
get_slice_tile(
@@ -1131,13 +1200,19 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
{
block_sync_lds();
gemm_1(
o_acc,
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::Q_PERTENSOR_KV_BLOCKSCALE)
{
tile_elementwise_inout(
[&v_descale](auto& o, auto& o0) { o += o0 * v_descale; }, o_acc, o_acc0);
}
} while(i_total_loops < num_total_loop);
// store lse
@@ -1255,6 +1330,9 @@ struct BlockFmhaBatchPrefillPipelineQRKSVSAsync
page_stride_k,
page_stride_v,
dropout,
nullptr,
nullptr,
1,
sink_v);
}
};