diff --git a/include/ck_tile/ops/fmha/block/block_attention_quant_scale_enum.hpp b/include/ck_tile/ops/fmha/block/block_attention_quant_scale_enum.hpp index 7e0f704bef..4b7fa1b1a6 100644 --- a/include/ck_tile/ops/fmha/block/block_attention_quant_scale_enum.hpp +++ b/include/ck_tile/ops/fmha/block/block_attention_quant_scale_enum.hpp @@ -13,6 +13,7 @@ enum class BlockAttentionQuantScaleEnum NO_SCALE = 0, PERTENSOR = 1, BLOCKSCALE, + Q_PERTENSOR_KV_BLOCKSCALE, }; template @@ -33,5 +34,10 @@ struct BlockAttentionQuantScaleEnumToStr +struct BlockAttentionQuantScaleEnumToStr +{ + static constexpr const char* name = "q_pertensor_kv_blockscale"; +}; } // namespace ck_tile diff --git a/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp index 86e1de3e9f..45ecf7a18b 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_batch_prefill_kernel.hpp @@ -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 @@ -255,9 +274,13 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel FmhaFwdEmptyKargs<0>>>, std::conditional_t>, std::conditional_t>, - std::conditional_t>, + std::conditional_t< + QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR, + FmhaFwdCommonQScaleKargs, + std::conditional_t>>, std::conditional_t>, std::conditional_t> { @@ -276,9 +299,13 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel FmhaFwdEmptyKargs<0>>>, std::conditional_t>, std::conditional_t>, - std::conditional_t>, + std::conditional_t< + QScaleEnum == BlockAttentionQuantScaleEnum::PERTENSOR, + FmhaFwdCommonQScaleKargs, + std::conditional_t>>, std::conditional_t>, std::conditional_t> { @@ -348,7 +375,12 @@ struct FmhaBatchPrefillWithPagedKVCacheKernel bool s_randval, std::variant, std::pair> 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> 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(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(kargs.k_descale_ptr) + + static_cast(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(kargs.v_descale_ptr) + + static_cast(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(kargs.k_descale_ptr) + + static_cast(i_nhead / kargs.nhead_ratio_qk) * + kargs.nhead_stride_k_descale + + static_cast(i_batch) * kargs.batch_stride_k_descale; + v_descale_ptr = reinterpret_cast(kargs.v_descale_ptr) + + static_cast(i_nhead / kargs.nhead_ratio_qk) * + kargs.nhead_stride_v_descale + + static_cast(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( @@ -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(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); } }(); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp index 48e8f75ae7..911ee2f295 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_batch_prefill_pipeline_qr_ks_vs_async.hpp @@ -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{}) + 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{}); update_v_offsets(number{}); @@ -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{}), 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{}), get_slice_tile( v_lds_window, sequence<(LdsSeq.at(number{})) * kN1, 0>{}, sequence<(LdsSeq.at(number{}) + 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); } };