diff --git a/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp b/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp index 4732009cde454e..f3ef31dff99ec3 100644 --- a/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp +++ b/src/common/transformations/tests/common_optimizations/sdpa_to_vlsdpa_test.cpp @@ -58,24 +58,16 @@ std::shared_ptr build_target_model(const string& mask_name) { k->set_friendly_name("k"); v->set_friendly_name("v"); - auto transpose_q = std::make_shared(q, Constant::create(element::i64, Shape{3}, {1, 0, 2})); - auto transpose_k = std::make_shared(k, Constant::create(element::i64, Shape{3}, {1, 0, 2})); - auto transpose_v = std::make_shared(v, Constant::create(element::i64, Shape{3}, {1, 0, 2})); - transpose_q->set_friendly_name("transpose_q"); - transpose_k->set_friendly_name("transpose_k"); - transpose_v->set_friendly_name("transpose_v"); - auto cuseq_mask = std::make_shared(element::i32, PartialShape{-1}); cuseq_mask->set_friendly_name(mask_name); cuseq_mask->get_output_tensor(0).set_names({mask_name}); + const std::vector order{1, 0, 2}; auto vlsdpa = - std::make_shared(OutputVector{transpose_q, transpose_k, transpose_v, cuseq_mask}); - - auto transpose_o = std::make_shared(vlsdpa, Constant::create(element::i64, Shape{3}, {1, 0, 2})); - transpose_o->set_friendly_name("transpose_o"); + std::make_shared(OutputVector{q, k, v, cuseq_mask}, order, order, order, order); + vlsdpa->set_friendly_name("transpose_o"); - return std::make_shared(OutputVector{transpose_o}, ParameterVector{q, k, v, cuseq_mask}); + return std::make_shared(OutputVector{vlsdpa}, ParameterVector{q, k, v, cuseq_mask}); } }; // namespace diff --git a/src/core/src/pass/sdpa_to_vlsdpa.cpp b/src/core/src/pass/sdpa_to_vlsdpa.cpp index 9988ed4bd2ff1c..8f5f90a205dc0d 100644 --- a/src/core/src/pass/sdpa_to_vlsdpa.cpp +++ b/src/core/src/pass/sdpa_to_vlsdpa.cpp @@ -9,7 +9,11 @@ #include #include "openvino/cc/pass/itt.hpp" +#include "openvino/core/graph_util.hpp" +#include "openvino/core/rt_info.hpp" +#include "openvino/op/constant.hpp" #include "openvino/op/scaled_dot_product_attention.hpp" +#include "openvino/op/transpose.hpp" #include "openvino/op/util/node_util.hpp" #include "openvino/pass/manager.hpp" #include "ov_ops/vl_sdpa.hpp" @@ -85,14 +89,74 @@ bool SDPAToVLSDPA::run_on_model(const std::shared_ptr& model) { const auto sdpa_consumers = sdpa->get_output_target_inputs(0); const auto new_args = sdpa->input_values(); - OutputVector inputs{new_args.at(0), new_args.at(1), new_args.at(2), cu_seqlens_param}; + + // Try to fuse the {1,0,2} Transposes surrounding SDPA into the VLSDPA op + // at op-creation time. Doing it here (rather than as a later plugin-side + // matcher on `Transpose->VLSDPA->Transpose`) is race-free w.r.t. common + // optimizations that may convert Transpose->Reshape when one of the + // permuted dims degenerates to size 1 (e.g. num_head==1). Constraints + // mirror `VLSDPA::validate_and_infer_types`: all four orders must be + // equal and, when non-empty, equal to {1,0,2}. + const std::vector expected_order{1, 0, 2}; + auto try_peel_input_transpose = + [&expected_order](ov::Output& value) -> std::shared_ptr { + auto transpose = ov::as_type_ptr(value.get_node_shared_ptr()); + if (!transpose) + return nullptr; + if (transpose->get_output_target_inputs(0).size() != 1) + return nullptr; + auto order_const = ov::as_type_ptr(transpose->input_value(1).get_node_shared_ptr()); + if (!order_const || ov::shape_size(order_const->get_shape()) != expected_order.size()) + return nullptr; + if (order_const->cast_vector() != expected_order) + return nullptr; + value = transpose->input_value(0); + return transpose; + }; + + ov::Output q_in = new_args.at(0); + ov::Output k_in = new_args.at(1); + ov::Output v_in = new_args.at(2); + auto tq = try_peel_input_transpose(q_in); + auto tk = try_peel_input_transpose(k_in); + auto tv = try_peel_input_transpose(v_in); + + std::shared_ptr to; + if (sdpa_consumers.size() == 1) { + auto candidate_out = + ov::as_type_ptr(sdpa_consumers.begin()->get_node()->shared_from_this()); + if (candidate_out) { + auto out_order_const = + ov::as_type_ptr(candidate_out->input_value(1).get_node_shared_ptr()); + if (out_order_const && ov::shape_size(out_order_const->get_shape()) == expected_order.size() && + out_order_const->cast_vector() == expected_order) { + to = candidate_out; + } + } + } + + const bool fuse_transposes = tq && tk && tv && to; std::shared_ptr vl_sdpa; - vl_sdpa = std::make_shared(inputs); - vl_sdpa->set_friendly_name(sdpa->get_friendly_name()); + if (fuse_transposes) { + OutputVector inputs{q_in, k_in, v_in, cu_seqlens_param}; + vl_sdpa = std::make_shared(inputs, + expected_order, + expected_order, + expected_order, + expected_order); + vl_sdpa->set_friendly_name(to->get_friendly_name()); + ov::copy_runtime_info({tq, tk, tv, sdpa->shared_from_this(), to}, vl_sdpa); + ov::replace_node(to, vl_sdpa); + } else { + OutputVector inputs{new_args.at(0), new_args.at(1), new_args.at(2), cu_seqlens_param}; + vl_sdpa = std::make_shared(inputs); + vl_sdpa->set_friendly_name(sdpa->get_friendly_name()); + ov::copy_runtime_info(sdpa->shared_from_this(), vl_sdpa); - for (auto& consumer : sdpa_consumers) - consumer.replace_source_output(vl_sdpa); + for (auto& consumer : sdpa_consumers) + consumer.replace_source_output(vl_sdpa); + } } } } diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm b/src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm index 87ac3c56222d2e..75e15ba1d4b081 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm +++ b/src/plugins/intel_gpu/src/graph/impls/cm/cm_sdpa_vlen.cm @@ -15,7 +15,7 @@ * limitations under the License. *******************************************************************************/ namespace KERNEL_NAME { - + #include "cm_sdpa_common.hpp" #ifdef CM_HAS_LSC_UNTYPED_2D @@ -36,10 +36,32 @@ _GENX_MAIN_ void KERNEL_NAME( SurfaceIndex query [[type("buffer_t")]], SurfaceIndex key [[type("buffer_t")]], SurfaceIndex value [[type("buffer_t")]], - SurfaceIndex output [[type("buffer_t")]], + SurfaceIndex output [[type("buffer_t")]], #endif int* cu_seqlens [[type("svmptr_t")]], - int need_wg_mapping + int need_wg_mapping, + // Memory-layout contract for Q/K/V inputs (must be preserved by upstream passes): + // Only the token (outer) axis is allowed to be discontinuous. The kernel treats + // each input as a 3-D tensor [tokens, heads, head_size] and hardcodes the two + // inner strides: + // - per-head stride == head_size (heads contiguous within a token) + // - per-element stride == 1 (head_size dim is contiguous) + // Only the outer per-token stride is layout-driven via {q,k,v}_token_pitch below. + // If prepare_buffer_fusing ever propagates padding onto the head or head_size axis + // of a vl_sdpa input layout, this kernel will silently address the wrong memory. + // + // Element-count offsets from the SVM base to the slice start of each input, taken + // from input_layout.get_linear_offset(). Non-zero when an in-place crop (e.g. the + // TransposeSplitMatcher axis=1 pattern) leaves the SVM base pointing at the packed + // parent buffer and shifts the logical view via layout padding on the token axis. + int token_offset_q, + int token_offset_k, + int token_offset_v, + // Per-input token strides in elements, from input_layout.get_pitches()[0]. Allows + // mixed layouts (e.g. Q/K contiguous post-RoPE, V still viewing the padded parent). + int q_token_pitch, + int k_token_pitch, + int v_token_pitch ) { constexpr int num_heads = CMFLA_NUM_HEADS; constexpr int head_size = CMFLA_HEAD_SIZE; @@ -107,8 +129,12 @@ _GENX_MAIN_ void KERNEL_NAME( kv_start = cu_seqlens[wg_id]; kv_seq_len = cu_seqlens[wg_id + 1] - kv_start; } - uint qo_offset = (q_start*num_heads + h)*head_size; - uint kv_offset = (kv_start*num_kv_heads + hkv)*head_size; + // Input offsets are computed with per-input runtime token pitches. + // Output is always contiguous [q_len, num_heads, head_size]. + uint q_in_offset = q_start * q_token_pitch + h * head_size; + uint k_in_offset = kv_start * k_token_pitch + hkv * head_size; + uint v_in_offset = kv_start * v_token_pitch + hkv * head_size; + uint o_offset = (q_start * num_heads + h) * head_size; #if USE_LSC == 1 #if USE_LSC_PREFETCH @@ -118,10 +144,13 @@ _GENX_MAIN_ void KERNEL_NAME( kv_seq_len, q_len, kv_seq_len, //kv_len, - reinterpret_cast(query + qo_offset), - reinterpret_cast(key + kv_offset), - reinterpret_cast(value + kv_offset), - reinterpret_cast(output + qo_offset)); + q_token_pitch * sizeof(half), + k_token_pitch * sizeof(half), + v_token_pitch * sizeof(half), + reinterpret_cast((query + token_offset_q) + q_in_offset), + reinterpret_cast((key + token_offset_k) + k_in_offset), + reinterpret_cast((value + token_offset_v) + v_in_offset), + reinterpret_cast(output + o_offset)); #else sdpa_kernel_lsc( slm_K, @@ -132,10 +161,13 @@ _GENX_MAIN_ void KERNEL_NAME( kv_seq_len, //kv_stop, q_len, //q_len, kv_seq_len, //kv_len, - reinterpret_cast(query + qo_offset), - reinterpret_cast(key + kv_offset), - reinterpret_cast(value + kv_offset), - reinterpret_cast(output + qo_offset)); + q_token_pitch * sizeof(half), + k_token_pitch * sizeof(half), + v_token_pitch * sizeof(half), + reinterpret_cast((query + token_offset_q) + q_in_offset), + reinterpret_cast((key + token_offset_k) + k_in_offset), + reinterpret_cast((value + token_offset_v) + v_in_offset), + reinterpret_cast(output + o_offset)); #endif #else sdpa_kernel( @@ -151,10 +183,13 @@ _GENX_MAIN_ void KERNEL_NAME( key, value, output, - qo_offset * sizeof(half), - kv_offset * sizeof(half), - kv_offset * sizeof(half), - qo_offset * sizeof(half) + q_token_pitch * sizeof(half), + k_token_pitch * sizeof(half), + v_token_pitch * sizeof(half), + (token_offset_q + q_in_offset) * sizeof(half), + (token_offset_k + k_in_offset) * sizeof(half), + (token_offset_v + v_in_offset) * sizeof(half), + o_offset * sizeof(half) ); #endif } diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp index 9ab9b86930e235..7dc8c6ca0d27ab 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_attention_common.hpp @@ -7,8 +7,6 @@ //# CM-compiler is C++17 static_assert(__cplusplus >= 201703L); -//# static_assert(__cplusplus >= 202002L); -//# static_assert(__cplusplus >= 202302L); #define SystolicDepth 8 #define RepeatCount 8 @@ -23,13 +21,41 @@ static_assert(__cplusplus >= 201703L); #define kv_step REG_K #define q_step REG_N +// Q pre-scale math constants: +// scale_factor = 1 / sqrt(head_size) (provided by CMFLA_SCALE_FACTOR at JIT time) +// log2e = log2(e) (fixed) constexpr float scale_factor = CMFLA_SCALE_FACTOR; +constexpr float log2e = 1.4426950408889634f; + +// JIT const: chooses the domain the softmax score St = K @ Q^T lives in. +// 1 => Q is pre-multiplied by log2e so the online softmax can call cm_exp (== exp2) +// directly (Item 13 in OPTIMIZE_PLAN, drops one *log2e off the softmax critical path). +// 0 => Q keeps the natural-log domain and the softmax multiplies by log2e internally. +// Set at JIT time (typically via -DCMFLA_Q_SCALED_BY_LOG2=... from the host Python script); +// defaults to 1. The value MUST match how the kernel actually pre-scales Q -- both the Q +// load (via q_prescale) and the softmax templates read this macro, so keeping them in +// sync is enough to prevent log2e from being applied zero or two times. +#ifndef CMFLA_Q_SCALED_BY_LOG2 +#define CMFLA_Q_SCALED_BY_LOG2 1 +#endif + +// Single Q pre-scale constant used by every SDPA/PA kernel at Q load time. +constexpr float q_prescale = CMFLA_Q_SCALED_BY_LOG2 ? (scale_factor * log2e) : scale_factor; + +// JIT const: 1 => use the balanced binary-tree reduction (depth log2(rows)) in the online +// softmax; requires rows to be a power of two. +// 0 => use the linear-chain reduction (default). +// Applies uniformly to the SDPA (flashattn) and PA (pageatten) kernel families via the +// `cm_online_softmax_update` dispatch macro defined below. +#ifndef CMFLA_USE_TREE_SOFTMAX +#define CMFLA_USE_TREE_SOFTMAX 0 +#endif static_assert(q_step == 16 || q_step == 8); static_assert(kv_step == 16); static_assert(CM_HAS_DPAS); -#define DEBUG_SHOW 1 +#define DEBUG_SHOW 0 #if !DEBUG_SHOW template void show(const matrix mat, bool isfloat=true) { @@ -217,9 +243,6 @@ inline matrix ugemm_KQ(uint slm_K, matrix_ref Kmat; cm_slm_block_read(slm_K, GENX_NONE, slm_offset, Kmat.format()); - // if (cm_local_id(2) == 3 && cm_group_id(2) == 0) { - // show(Kmat.format()); - // } #pragma unroll for(int k = 0; k < num_K; k++) St2.row(k) = cm_dpas(0, Qt[0].format(), Kmat[k].format()); @@ -244,16 +267,12 @@ inline void ugemm_PV0(uint slm_V, matrix_ref P, matrix_ref Vmat; cm_slm_block_read(slm_V, GENX_NONE, slm_offset + REG_K*k*sizeof(half), Vmat.format()); - // if (cm_local_id(2) == 0 && cm_group_id(2) == 0) { - // show(Vmat.format()); - // } #pragma unroll for(int p = 0; p < num_P_tiles; p++) { rO[ri + p] = cm_dpas( 0, Vmat.format(), P2.row(p).format()); - //show(rO[ri + p].format()); } } } @@ -268,9 +287,6 @@ inline void ugemm_PV1(uint slm_V, matrix_ref P, vector_ref Vmat; cm_slm_block_read(slm_V, GENX_NONE, slm_offset + REG_K*k*sizeof(half), Vmat.format()); - // if (cm_local_id(2) == 0 && cm_group_id(2) == 0) { - // show(Vmat.format()); - // } #pragma unroll for(int p = 0; p < num_P_tiles; p++) { auto cO = rO[ri + p].format(); @@ -279,7 +295,6 @@ inline void ugemm_PV1(uint slm_V, matrix_ref P, vector_ref(cO.row(r), max_comp[r + p*REG_M]); } - //show(rO[ri].format()); #pragma unroll for(int p = 0; p < num_P_tiles; p++) { @@ -287,12 +302,11 @@ inline void ugemm_PV1(uint slm_V, matrix_ref P, vector_ref(), Vmat.format(), P2.row(p).format()); - //if (kv_pos == args_verbose) show(rO[ri + p].format()); } - // if (kv_pos == args_verbose) show(cur_O.format()); } } +// Online-softmax update (linear-chain max/sum reduction). template vector online_softmax_update(matrix_ref St, vector_ref cur_max, vector_ref cur_sum) { vector new_max_t; @@ -301,86 +315,109 @@ vector online_softmax_update(matrix_ref St, vector_r new_max_t = cm_max(new_max_t, cur_max); // Pt = torch.exp(St - new_max) - constexpr float log2e = 1.4426950408889634f; +#if CMFLA_Q_SCALED_BY_LOG2 + for(int r = 0; r < St.n_rows(); r++) St[r] = cm_exp(St[r] - new_max_t); +#else for(int r = 0; r < St.n_rows(); r++) St[r] = cm_exp((St[r] - new_max_t)*log2e); +#endif vector row_sum_t; row_sum_t = cm_add(St[0], St[1]); for(int r = 2; r < St.n_rows(); r++) row_sum_t = cm_add(row_sum_t, St[r]); vector max_comp; +#if CMFLA_Q_SCALED_BY_LOG2 + max_comp = cm_exp(cur_max - new_max_t); +#else max_comp = cm_exp((cur_max - new_max_t)*log2e); +#endif cur_sum = cm_mul(cur_sum, max_comp); cur_sum = cm_add(cur_sum, row_sum_t); cur_max = new_max_t; return max_comp; } -#ifdef CM_HAS_LSC_UNTYPED_2D - #define cm_load_normal cm_load - #define cm_load_transpose cm_load - #define cm_load_vnni cm_load - #define cm_store_normal cm_store - // // simulation of LSC API using SVM API - // template - // inline void cm_load_normal(vector_ref Res, const lsc::block_2d_desc &Desc, int16_t Pred = 1) { - // static_assert(NBlocks == 1); - // auto pitch = Desc.get_pitch() + 1; - // auto base = reinterpret_cast(Desc.get_base() + Desc.get_block_y()*pitch + Desc.get_block_x() * sizeof(T)); - // #pragma unroll - // for(int i = 0; i < BlockH; i++) { - // cm_svm_block_read(base + i * pitch, Res.select(i*BlockW)); - // } - // } - - // template - // inline void cm_load_transpose(vector_ref Res, const lsc::block_2d_desc &Desc, int16_t Pred = 1) { - // static_assert(NBlocks == 1); - // auto pitch = Desc.get_pitch() + 1; - // auto base = reinterpret_cast(Desc.get_base() + Desc.get_block_y()*pitch + Desc.get_block_x() * sizeof(T)); - // matrix temp; - // #pragma unroll - // for(int i = 0; i < BlockH; i++) { - // cm_svm_block_read(base + i * pitch, temp[i]); - // } - // Transpose2DMatrix(temp, Res.format()); - // } - - // // in VNNI case, NBlocks is increasing along X dimension (increase cache-line usage) - // template - // inline void cm_load_vnni(vector_ref Res, const lsc::block_2d_desc &Desc, int16_t Pred = 1) { - // static_assert(NBlocks == 1 || NBlocks == 2); - // // each block must be a full XMX B matrix - // static_assert(BlockH == REG_K); - // static_assert(BlockW == REG_N); - // auto pitch = Desc.get_pitch() + 1; - // auto base = reinterpret_cast(Desc.get_base() + Desc.get_block_y()*pitch + Desc.get_block_x() * sizeof(T)); - // matrix temp; - // #pragma unroll - // for(int i = 0; i < BlockH; i++) { - // cm_svm_block_read(base + i * pitch, temp[i]); - // } - - // auto out_vnni = Res.format(); - // #pragma unroll - // for(int i = 0; i < NBlocks; i ++) { - // out_vnni.select(i*(BlockH/2), 0) = temp.select(0, i*BlockW); - // out_vnni.select(i*(BlockH/2), 1) = temp.select(1, i*BlockW); - // } - // } - - // template - // inline void cm_store_normal(const lsc::block_2d_desc &Desc, vector_ref Res) { - // static_assert(NBlocks == 1); - // auto pitch = Desc.get_pitch() + 1; - // auto base = reinterpret_cast(Desc.get_base() + Desc.get_block_y()*pitch + Desc.get_block_x() * sizeof(T)); - // #pragma unroll - // for(int i = 0; i < BlockH; i++) { - // cm_svm_block_write(base + i * pitch, Res.select(i*BlockW)); - // } - // } +// Tree-reduction variant: max/sum reductions use a balanced binary tree (depth log2(rows)) +// instead of a linear chain (depth rows-1), shortening the loop-carried dependency chain. +// Requires rows to be a power of two; the PA kv_step=16 and KV_BLK*kv_step satisfy this. +template +CM_INLINE vector online_softmax_update_tree(matrix_ref St, + vector_ref cur_max, + vector_ref cur_sum) { + static_assert((rows & (rows - 1)) == 0, "tree reduction needs power-of-two rows"); + vector new_max_t; + { + matrix 1 ? rows/2 : 1), cols> t; + #pragma unroll + for (int r = 0; r < rows/2; r++) t.row(r) = cm_max(St[r], St[r + rows/2]); + #pragma unroll + for (int stride = rows/4; stride > 0; stride >>= 1) + #pragma unroll + for (int r = 0; r < stride; r++) + t.row(r) = cm_max(t.row(r), t.row(r + stride)); + new_max_t = cm_max(t.row(0), cur_max); + } +#if CMFLA_Q_SCALED_BY_LOG2 + #pragma unroll + for (int r = 0; r < rows; r++) St[r] = cm_exp(St[r] - new_max_t); +#else + #pragma unroll + for (int r = 0; r < rows; r++) St[r] = cm_exp((St[r] - new_max_t) * log2e); #endif + vector row_sum_t; + { + matrix 1 ? rows/2 : 1), cols> t; + #pragma unroll + for (int r = 0; r < rows/2; r++) t.row(r) = cm_add(St[r], St[r + rows/2]); + #pragma unroll + for (int stride = rows/4; stride > 0; stride >>= 1) + #pragma unroll + for (int r = 0; r < stride; r++) + t.row(r) = cm_add(t.row(r), t.row(r + stride)); + row_sum_t = t.row(0); + } + + vector max_comp; +#if CMFLA_Q_SCALED_BY_LOG2 + max_comp = cm_exp(cur_max - new_max_t); +#else + max_comp = cm_exp((cur_max - new_max_t) * log2e); +#endif + cur_sum = cm_mul(cur_sum, max_comp); + cur_sum = cm_add(cur_sum, row_sum_t); + cur_max = new_max_t; + return max_comp; +} + +// Dispatch macro used by every SDPA/PA kernel; selects linear vs tree reduction. +#if CMFLA_USE_TREE_SOFTMAX +#define cm_online_softmax_update(St, cur_max, cur_sum) online_softmax_update_tree(St, cur_max, cur_sum) +#else +#define cm_online_softmax_update(St, cur_max, cur_sum) online_softmax_update(St, cur_max, cur_sum) +#endif + +// Transpose a float score tile (kv x q) into a half P tile (q x kv) for the P@V matmul. +// Casting float->half first (one vectorized cm_mul-free copy) lets the shuffle network +// run at half width -- roughly halving the mov count vs transposing the float tile +// directly through Transpose_*x*. +// +// Two overloads cover both Xe generations (kv_step is always 16; q_step = REG_N): +// * [16,16] -> [16,16] for Xe2 (q_step = 16) +// * [16, 8] -> [ 8,16] for Xe1 (q_step = 8) +CM_INLINE void transpose_St_to_P_half(matrix_ref St, matrix_ref P) { + matrix Sh; + #pragma unroll + for (int r = 0; r < 16; r++) Sh.row(r) = St.row(r); + Transpose_16x16(Sh.select<16,1,16,1>(0,0), P); +} +CM_INLINE void transpose_St_to_P_half(matrix_ref St, matrix_ref P) { + matrix Sh; + #pragma unroll + for (int r = 0; r < 16; r++) Sh.row(r) = St.row(r); + Transpose2DMatrix(Sh.select<16,1,8,1>(0,0), P); +} + //=============================================================================================== template constexpr void apply_causal_mask(matrix_ref St) { diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe1.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe1.hpp index ad98fc979928e8..6eec4e4a8f8635 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe1.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe1.hpp @@ -97,7 +97,7 @@ void pa_lsc_u8( CacheHint::Cached, CacheHint::Cached>(q_gather, gather_offsets, gather_pred); rQ[ri].format() = gathered; - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } #if KV_CACHE_COMPRESSION == 1 @@ -273,7 +273,6 @@ void pa_lsc_u8( cm_fence(CM_LOCAL_BARRIER); cm_sbarrier(0); - //if (kv_pos > 1024000) if (kv_pos + kv_step < kv_stop) cm_sbarrier(1); load_slm_KV(kv_pos + kv_step*2); @@ -300,7 +299,7 @@ void pa_lsc_u8( // LSC ensures no overflow-access, but mask off k-tails attn-score is still required for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; } - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; Transpose2DMatrix(St, P); @@ -480,7 +479,7 @@ void pa_kernel_lsc_prefetch_f16( CacheHint::Cached, CacheHint::Cached>(q_gather, gather_offsets, gather_pred); rQ[ri].format() = gathered; - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } constexpr int blk_stride = CMFLA_NUM_KV_HEADS * CMFLA_HEAD_SIZE*CMPA_BLOCK_SZ; @@ -568,8 +567,7 @@ void pa_kernel_lsc_prefetch_f16( for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; } - // show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; Transpose2DMatrix(St, P); diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp index cc88251d6ed00c..bc60e3077ed5a6 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_pa_xe2.hpp @@ -126,7 +126,7 @@ void pa_lsc_u8( #pragma unroll for (int k = 0, ri = 0; k < process_head_size / 2; k += REG_K / 2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(worker_offset / 2 + k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } @@ -398,10 +398,10 @@ void pa_lsc_u8( } int kv_tokens = kv_stop - kv_pos; for (int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; - Transpose2DMatrix(St, P); + transpose_St_to_P_half(St, P); // Each worker reads its chunk of V from SLM uint slm_V_worker_lo_offset = worker_offset * REG_K * sizeof(half); @@ -632,10 +632,10 @@ void pa_lsc_u8( } int kv_tokens = kv_stop - kv_pos; for (int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; - Transpose2DMatrix(St, P); + transpose_St_to_P_half(St, P); // Each worker reads its chunk of V from SLM uint slm_V_worker_lo_offset = worker_offset * REG_K * sizeof(half); @@ -847,7 +847,7 @@ void pa_kernel_lsc_prefetch_f16( #pragma unroll for(int k = 0, ri = 0; k < process_head_size/2; k += REG_K/2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(worker_offset / 2 + k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } @@ -967,11 +967,10 @@ void pa_kernel_lsc_prefetch_f16( int kv_tokens = kv_stop - kv_pos; // LSC ensures no overflow-access, but mask off k-tails attn-score is still required for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; - Transpose2DMatrix(St, P); + transpose_St_to_P_half(St, P); prefetch_V.set_base_ptr((reinterpret_cast(v_cache_base)+prefetch_block_id*blk_stride)); prefetch_V.set_block_y((prefetch_kv_pos + wg_local_id) % CMPA_BLOCK_SZ); @@ -1271,11 +1270,10 @@ void pa_kernel_lsc_prefetch_f16( int kv_tokens = kv_stop - kv_pos; // LSC ensures no overflow-access, but mask off k-tails attn-score is still required for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; - Transpose2DMatrix(St, P); + transpose_St_to_P_half(St, P); prefetch_V.set_base_ptr((reinterpret_cast(v_cache_base)+prefetch_block_id*blk_stride)); prefetch_V.set_block_y((prefetch_kv_pos + wg_local_id) % CMPA_BLOCK_SZ); diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp index d9c368406014c8..78fa3f98aa039b 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/include/cm_sdpa_common.hpp @@ -1,270 +1,19 @@ -/******************************************************************************* - * Copyright (c) 2018-2026 Intel Corporation - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - *******************************************************************************/ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + #include "cm_attention_common.hpp" -#ifdef CM_HAS_LSC_UNTYPED_2D -//@prefetch_u8 would have duplicated decompress perf issue. comments out for now. -// template -// void sdpa_kernel_lsc_prefetch_u8( -// int wg_local_id, -// int q_start, -// int kv_stop, // -// int q_len, //q_step -// int kv_len, //not used for now -// svmptr_t q_base [[type("svmptr_t")]], -// svmptr_t k_cache_base [[type("svmptr_t")]], -// svmptr_t v_cache_base [[type("svmptr_t")]], -// svmptr_t o_base [[type("svmptr_t")]], -// int32_t past_lens, -// int32_t* block_indices [[type("svmptr_t")]]) { -// constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); -// constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; -// //[block_num, kv_heads, block_size, head_size] -// constexpr uint kv_pitch = head_size * sizeof(uint8_t); - -// vector cur_max; -// vector cur_sum; - -// cur_max = -3e38f; -// cur_sum = 0; -// constexpr int num_P_tiles = REG_N / REG_M; -// matrix rQ; -// matrix rO; - -// auto q_tokens_left = q_len;// - q_start; -// static_assert(q_step == REG_N); -// static_assert(kv_step == REG_K); - -// if (q_tokens_left < 0) q_tokens_left = 0; -// if (q_tokens_left > q_step) q_tokens_left = q_step; - -// if (q_tokens_left > 0) { -// lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); -// #pragma unroll -// for(int k = 0, ri = 0; k < head_size/2; k += REG_K/2, ri++) { -// cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); -// rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); -// } -// } - -// lsc::block_2d_desc b2dKV(k_cache_base, CMPA_BLOCK_SZ - 1, head_size*sizeof(uint8_t) - 1, kv_pitch - 1, 0, 0); - -// static_assert(wg_local_size == 16); -// lsc::block_2d_desc b2dKV_prefetch(k_cache_base, CMPA_BLOCK_SZ - 1, head_size*sizeof(uint8_t) - 1, kv_pitch - 1, 0, 0); -// // constexpr int blk_stride = CMFLA_NUM_KV_HEADS * CMFLA_HEAD_SIZE * CMPA_BLOCK_SZ; -// constexpr int quan_blk_stride = CMFLA_NUM_KV_HEADS * (CMFLA_HEAD_SIZE+4) * CMPA_BLOCK_SZ * sizeof(uint8_t); - - - -// int causal_left = q_start+past_lens; -// for(int kv_pos = 0; kv_pos < kv_stop; kv_pos += kv_step) { -// auto cur_block_id = block_indices[kv_pos / CMPA_BLOCK_SZ]; -// //For the last step, duplicate prefetch here. -// uint32_t prefetch_kv_pos = (kv_pos+kv_step) >= kv_stop ? kv_pos : (kv_pos+kv_step); -// auto prefetch_block_id = block_indices[prefetch_kv_pos / CMPA_BLOCK_SZ]; -// uint32_t dscale_offset = cur_block_id*quan_blk_stride + CMPA_BLOCK_SZ * head_size * sizeof(uint8_t) + kv_pos%CMPA_BLOCK_SZ*sizeof(half); - -// vector dscale; -// vector zp; -// cm_svm_block_read(reinterpret_cast( k_cache_base + dscale_offset), dscale); -// cm_svm_block_read(reinterpret_cast( k_cache_base + dscale_offset + CMPA_BLOCK_SZ*sizeof(half)), zp); - -// // if (cm_local_id(2) == 0 && cm_group_id(2) == 0) { -// // show(dscale.format()); -// // } -// //# St = k @ Qt - -// matrix St; // = ugemm_KQ(slm_K, rQ, slm_offset); -// { -// constexpr int num_K = kv_step/REG_M; -// auto St2 = St.format(); -// matrix Kmat; -// auto quan_Kmat = Kmat.format().row(1).format(); -// auto dq_Kmat = Kmat.format(); -// //cm_slm_block_read(slm_K, GENX_NONE, slm_offset, Kmat.format()); - -// b2dKV_prefetch.set_base_ptr(reinterpret_cast(k_cache_base+prefetch_block_id*quan_blk_stride)); -// b2dKV_prefetch.set_block_y((prefetch_kv_pos + wg_local_id) % CMPA_BLOCK_SZ); -// cm_prefetch(b2dKV_prefetch.set_block_x(0)); - -// b2dKV.set_base_ptr(reinterpret_cast(k_cache_base+cur_block_id*quan_blk_stride)); -// b2dKV.set_block_y(kv_pos%CMPA_BLOCK_SZ); - -// cm_load(quan_Kmat.format(), b2dKV.set_block_x(0)); -// // if (cm_local_id(2) == 0 && cm_group_id(2) == 0) { -// // show(quan_Kmat.format(), false); -// // } -// #pragma unroll -// for(int r = 0; r < kv_step; r++) { -// dq_Kmat[r] = quan_Kmat[r] - zp[r]; -// dq_Kmat[r] = cm_mul(dq_Kmat[r], dscale[r]); -// } - -// #pragma unroll -// for(int k = 0; k < num_K; k++) -// St2.row(k) = cm_dpas( -// 0, -// rQ[0].format(), -// Kmat[k].format()); - -// #pragma unroll -// for(int ri = 1; ri < head_size/REG_K; ri++) { -// cm_prefetch(b2dKV_prefetch.set_block_x(ri*REG_K)); -// //cm_load(Kmat.format(), b2dKV.set_block_x(ri*REG_K)); -// cm_load(quan_Kmat.format(), b2dKV.set_block_x(ri*REG_K)); -// #pragma unroll -// for(int r = 0; r < kv_step; r++) { -// dq_Kmat[r] = quan_Kmat[r] - zp[r]; -// dq_Kmat[r] = cm_mul(dq_Kmat[r], dscale[r]); -// } -// #pragma unroll -// for(int k = 0; k < num_K; k++) { -// St2.row(k) = cm_dpas( -// St2.row(k), -// rQ[ri].format(), -// Kmat[k].format()); -// } -// } -// } -// if constexpr (use_causal_mask) { -// // since kv_step == q_step == 16, causal_left is n*kv_step -// if (causal_left == 0) { -// apply_causal_mask<1>(St); -// } else if (causal_left < 0) { -// St = -3.4e38f; -// } -// causal_left -= kv_step; -// } else { -// int kv_tokens = kv_stop - kv_pos; -// // LSC ensures no overflow-access, but mask off k-tails attn-score is still required -// for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; -// } - -// //show(St); -// auto max_comp = online_softmax_update(St, cur_max, cur_sum); - -// matrix P; -// Transpose2DMatrix(St, P); - -// b2dKV_prefetch.set_base_ptr(reinterpret_cast(v_cache_base+prefetch_block_id*quan_blk_stride)); -// b2dKV_prefetch.set_block_y((prefetch_kv_pos + wg_local_id) % CMPA_BLOCK_SZ); - -// b2dKV.set_base_ptr(reinterpret_cast(v_cache_base+cur_block_id*quan_blk_stride)); -// b2dKV.set_block_y(kv_pos%CMPA_BLOCK_SZ); - - -// cm_svm_block_read(reinterpret_cast(v_cache_base+dscale_offset), dscale); -// cm_svm_block_read(reinterpret_cast(v_cache_base+dscale_offset+CMPA_BLOCK_SZ*sizeof(half)), zp); - -// { -// matrix VmatVNNI2; -// matrix Vmat; -// auto quanVmat = Vmat.format().row(1).format(); -// int kv_tokens = kv_stop - kv_pos; -// if (kv_pos == 0) { -// // ugemm_PV0(slm_V, P, rO, slm_offset); -// auto P2 = P.format(); -// #pragma unroll -// for(int k = 0, ri = 0; k < head_size; k += REG_N, ri += num_P_tiles) { -// cm_prefetch(b2dKV_prefetch.set_block_x(k)); -// cm_load(quanVmat.format(), b2dKV.set_block_x(k)); -// #pragma unroll -// for(int r = 0; r < kv_step;r++) { -// Vmat[r] = quanVmat[r]-zp[r]; -// Vmat[r] = cm_mul(Vmat[r], dscale[r]); -// } -// for(int r = kv_step-1; r>=kv_tokens;r--) { -// Vmat[r] = 0; -// } - -// prepackAsVNNIWidth2(Vmat, VmatVNNI2); - -// #pragma unroll -// for(int p = 0; p < num_P_tiles; p++) { -// rO[ri + p] = cm_dpas( -// 0, -// VmatVNNI2.format(), -// P2.row(p).format()); -// } -// } -// } -// else { -// //ugemm_PV1(slm_V, P, max_comp, rO, slm_offset); -// auto P2 = P.format(); -// #pragma unroll -// for(int k = 0, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { -// cm_prefetch(b2dKV_prefetch.set_block_x(k)); -// cm_load(quanVmat.format(), b2dKV.set_block_x(k)); - -// #pragma unroll -// for(int r = 0; r < kv_step;r++) { -// Vmat[r] = quanVmat[r]-zp[r]; -// Vmat[r] = cm_mul(Vmat[r], dscale[r]); -// } -// for(int r = kv_step-1; r>=kv_tokens;r--) { -// Vmat[r] = 0; -// } - -// prepackAsVNNIWidth2(Vmat, VmatVNNI2); -// //# compensate cur_O -// // matrix rO; -// #pragma unroll -// for(int p = 0; p < num_P_tiles; p++) { -// auto cO = rO[ri + p].format(); -// #pragma unroll -// for(int r = 0; r < REG_M; r++) -// cO.row(r) = cm_mul(cO.row(r), max_comp[r + p*REG_M]); -// } - -// #pragma unroll -// for(int p = 0; p < num_P_tiles; p++) { -// rO[ri + p] = cm_dpas( -// rO[ri + p].format(), -// VmatVNNI2.format(), -// P2.row(p).format()); -// } -// } -// } -// } -// } -// if (q_tokens_left == 0) return; - -// //# save cur_O/cur_sum.transpose(0, 1) -// matrix cur_O_f16; -// cur_sum = cm_inv(cur_sum); - -// lsc::block_2d_desc b2dO(o_base, q_tokens_left - 1, head_size*sizeof(half) - 1, o_pitch - 1, 0, 0); - -// #pragma unroll -// for(int k = 0, ri=0; k < head_size; k += REG_N, ri += num_P_tiles) { -// #pragma unroll -// for(int p = 0; p < num_P_tiles; p++) { -// auto cO = rO[ri + p].format(); -// #pragma unroll -// for(int r = 0; r < cO.n_rows(); r++) { -// cur_O_f16[r + p*REG_M] = cm_mul(cO.row(r), cur_sum[r + p*REG_M]); - -// } -// } -// b2dO.set_block_x(k); -// cm_store(b2dO.set_block_y(0), cur_O_f16.format().row(0)); -// cm_store(b2dO.set_block_y(REG_M), cur_O_f16.format().row(1)); -// } -// } +// Number of kv sub-tiles (kv_step rows each) processed per online-softmax update in +// sdpa_kernel_lsc_prefetch. Larger amortizes the rO rescale + softmax bookkeeping over +// more DPAS work at the cost of GRF pressure. Overridable via -DCMFLA_KV_BLK=N. +#ifndef CMFLA_KV_BLK +#define CMFLA_KV_BLK 1 +#endif + +// online_softmax_update_tree and transpose_St_to_P_half are defined in cm_attention_common.hpp +#ifdef CM_HAS_LSC_UNTYPED_2D template void sdpa_kernel_lsc( uint slm_K, @@ -275,14 +24,15 @@ void sdpa_kernel_lsc( int kv_stop, int q_len, int kv_len, + uint q_pitch_bytes, + uint k_pitch_bytes, + uint v_pitch_bytes, svmptr_t q_base [[type("svmptr_t")]], svmptr_t k_base [[type("svmptr_t")]], svmptr_t v_base [[type("svmptr_t")]], svmptr_t o_base [[type("svmptr_t")]]) { constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); - constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; - constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); // round up head_size to multiple of 16 // block_2d_desc will automatically handle the tailing block constexpr int padded_head_size = (head_size + 16 - 1) / 16 * 16; @@ -294,6 +44,7 @@ void sdpa_kernel_lsc( constexpr int num_P_tiles = REG_N / REG_M; matrix rQ; matrix rO; + rO = 0.0f; // Zero the accumulator: the first softmax block scales rO by max_comp==0, and 0*NaN==NaN if the GRF holds stale NaN/Inf bits. auto q_tokens_left = q_len; static_assert(q_step == REG_N); @@ -303,16 +54,16 @@ void sdpa_kernel_lsc( if (q_tokens_left > q_step) q_tokens_left = q_step; if (q_tokens_left > 0) { - lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); + lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch_bytes - 1, 0, 0); #pragma unroll for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } - lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, k_pitch_bytes - 1, 0, 0); + lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, v_pitch_bytes - 1, 0, 0); int causal_left = q_start; @@ -348,8 +99,8 @@ void sdpa_kernel_lsc( cm_sbarrier(1); for(int kv_pos = 0; kv_pos < kv_stop; kv_pos += kv_step, - k_base += kv_step * kv_pitch, - v_base += kv_step * kv_pitch, + k_base += kv_step * k_pitch_bytes, + v_base += kv_step * v_pitch_bytes, slm_buff_id_read ++) { // load0, load1, signal1, @@ -365,7 +116,6 @@ void sdpa_kernel_lsc( cm_fence(CM_LOCAL_BARRIER); cm_sbarrier(0); - //if (kv_pos > 1024000) // for debugging if (kv_pos + kv_step < kv_stop) cm_sbarrier(1); @@ -390,8 +140,7 @@ void sdpa_kernel_lsc( for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; } - //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; Transpose2DMatrix(St, P); @@ -434,14 +183,15 @@ void sdpa_kernel_lsc_prefetch( int kv_stop, int q_len, int kv_len, + uint q_pitch_bytes, + uint k_pitch_bytes, + uint v_pitch_bytes, svmptr_t q_base [[type("svmptr_t")]], svmptr_t k_base [[type("svmptr_t")]], svmptr_t v_base [[type("svmptr_t")]], svmptr_t o_base [[type("svmptr_t")]]) { constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); - constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; - constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); // round up head_size to multiple of 16 // block_2d_desc will automatically handle the tailing block constexpr int padded_head_size = (head_size + 16 - 1) / 16 * 16; @@ -454,6 +204,7 @@ void sdpa_kernel_lsc_prefetch( constexpr int num_P_tiles = REG_N / REG_M; matrix rQ; matrix rO; + rO = 0.0f; auto q_tokens_left = q_len;// - q_start; static_assert(q_step == REG_N); @@ -463,119 +214,141 @@ void sdpa_kernel_lsc_prefetch( if (q_tokens_left > q_step) q_tokens_left = q_step; if (q_tokens_left > 0) { - lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch - 1, 0, 0); + lsc::block_2d_desc b2dQ(reinterpret_cast(q_base), q_tokens_left - 1, head_size*sizeof(half) - 1, q_pitch_bytes - 1, 0, 0); #pragma unroll for(int k = 0, ri = 0; k < padded_head_size/2; k += REG_K/2, ri++) { cm_load(rQ[ri].format(), b2dQ.set_block_x(k)); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } - lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + lsc::block_2d_desc b2dK(k_base, kv_stop - 1, head_size*sizeof(half) - 1, k_pitch_bytes - 1, 0, 0); + lsc::block_2d_desc b2dV(v_base, kv_stop - 1, head_size*sizeof(half) - 1, v_pitch_bytes - 1, 0, 0); static_assert(wg_local_size == 16); - lsc::block_2d_desc prefetch_K(k_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); - lsc::block_2d_desc prefetch_V(v_base, kv_stop - 1, head_size*sizeof(half) - 1, kv_pitch - 1, 0, 0); + lsc::block_2d_desc prefetch_K(k_base, kv_stop - 1, head_size*sizeof(half) - 1, k_pitch_bytes - 1, 0, 0); + lsc::block_2d_desc prefetch_V(v_base, kv_stop - 1, head_size*sizeof(half) - 1, v_pitch_bytes - 1, 0, 0); int causal_left = q_start; - for(int kv_pos = 0; kv_pos < kv_stop; kv_pos += kv_step, - k_base += kv_step * kv_pitch, - v_base += kv_step * kv_pitch) { - //# St = k @ Qt - matrix St; // = ugemm_KQ(slm_K, rQ, slm_offset); - { - constexpr int num_K = kv_step/REG_M; - auto St2 = St.format(); + // KV-blocking factor: number of kv tiles (each kv_step rows) processed per online- + // softmax update. The per-iteration rO rescale (REG_M*num_rO_tiles muls) and the + // softmax bookkeeping amortize over KV_BLK tiles, cutting the ALU-bound consume cost. + // The K@Q^T and the St->P transpose are still per sub-tile (they don't amortize). + constexpr int KV_BLK = CMFLA_KV_BLK; + constexpr int num_K = kv_step/REG_M; + constexpr int BLK_ROWS = KV_BLK*kv_step; + + // Warm up first K and V tiles before the loop so kv_pos=0 finds them in cache. + prefetch_K.set_block_y(wg_local_id); + prefetch_V.set_block_y(wg_local_id); + cm_prefetch(prefetch_K.set_block_x(0)); + cm_prefetch(prefetch_V.set_block_x(0)); + #pragma unroll + for(int ri = 1; ri < padded_head_size/REG_K; ri++) { + cm_prefetch(prefetch_K.set_block_x(ri*REG_K)); + cm_prefetch(prefetch_V.set_block_x(ri*REG_N)); + } + for(int kv_base = 0; kv_base < kv_stop; kv_base += BLK_ROWS) { + //# St = K @ Q^T for KV_BLK stacked sub-tiles -> [KV_BLK*kv_step, q_step] + matrix St; + auto St2 = St.format(); + #pragma unroll + for(int b = 0; b < KV_BLK; b++) { + int kv_pos = kv_base + b*kv_step; matrix Kmat; - //cm_slm_block_read(slm_K, GENX_NONE, slm_offset, Kmat.format()); - prefetch_K.set_block_y(wg_local_id + kv_pos + kv_step); + // Prefetch K of the matching tile one block ahead and V of this tile, so V + // is warm by the PV phase and the next block's K is warm by its K phase. + prefetch_K.set_block_y(wg_local_id + kv_pos + BLK_ROWS); + prefetch_V.set_block_y(wg_local_id + kv_pos); cm_prefetch(prefetch_K.set_block_x(0)); + cm_prefetch(prefetch_V.set_block_x(0)); b2dK.set_block_y(kv_pos); cm_load(Kmat.format(), b2dK.set_block_x(0)); #pragma unroll for(int k = 0; k < num_K; k++) - St2.row(k) = cm_dpas( + St2.row(b*num_K + k) = cm_dpas( 0, rQ[0].format(), Kmat[k].format()); #pragma unroll for(int ri = 1; ri < padded_head_size/REG_K; ri++) { - //cm_slm_block_read(slm_K, GENX_NONE, slm_offset + ri * Kmat.n_elems() * sizeof(half), Kmat.format()); cm_prefetch(prefetch_K.set_block_x(ri*REG_K)); + cm_prefetch(prefetch_V.set_block_x(ri*REG_N)); cm_load(Kmat.format(), b2dK.set_block_x(ri*REG_K)); #pragma unroll for(int k = 0; k < num_K; k++) { - St2.row(k) = cm_dpas( - St2.row(k), + St2.row(b*num_K + k) = cm_dpas( + St2.row(b*num_K + k), rQ[ri].format(), Kmat[k].format()); } } } + + // ---- mask per sub-tile (causal or kv-tail) ---- if constexpr (use_causal_mask) { - // since kv_step == q_step == 16, causal_left is n*kv_step - if (causal_left == 0) { - apply_causal_mask<1>(St); - } else if (causal_left < 0) { - St = -3.4e38f; + #pragma unroll + for(int b = 0; b < KV_BLK; b++) { + auto Stb = St.select(b*kv_step, 0); + int cl = causal_left - b*kv_step; + if (cl == 0) { + apply_causal_mask<1>(Stb); + } else if (cl < 0) { + Stb = -3.4e38f; + } } - causal_left -= kv_step; + causal_left -= BLK_ROWS; } else { - int kv_tokens = kv_stop - kv_pos; - // LSC ensures no overflow-access, but mask off k-tails attn-score is still required - for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; + int kv_tokens = kv_stop - kv_base; + // LSC ensures no overflow-access, but mask off k-tail attn-score is still required + for(int p = kv_tokens; p < BLK_ROWS; p++) St[p] = -3.4e38f; } - //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); - - matrix P; - Transpose2DMatrix(St, P); - - b2dV.set_block_y(kv_pos); - prefetch_V.set_block_y(wg_local_id +kv_pos + kv_step); - if (kv_pos == 0) { - // ugemm_PV0(slm_V, P, rO, slm_offset); - auto P2 = P.format(); - #pragma unroll - for(int k = 0, ri = 0; k < padded_head_size; k += REG_N, ri += num_P_tiles) { - matrix Vmat; - cm_prefetch(prefetch_V.set_block_x(k)); - cm_load(Vmat.format(), b2dV.set_block_x(k)); + // ---- one online-softmax update over the whole block ---- + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); + + // ---- rescale rO (skip first iter only when KV blocking amortizes the branch) ---- + // The kv_base=0 rescale is mathematically redundant for every head size: max_comp=0 + // and rO is still zero. The skip, however, adds a runtime kv_base branch to every + // outer KV block. Enable it only for KV_BLK>=2, where fewer outer iterations make + // the branch cost small enough to be profitable; KV_BLK=1 paths such as Omni HD=72 + // keep the straight-line rescale because forcing the branch measured as a regression. + if constexpr (KV_BLK >= 2) { + if (kv_base > 0) { #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - rO[ri + p] = cm_dpas( - 0, - Vmat.format(), - P2.row(p).format()); + for(int t = 0; t < padded_head_size/REG_N*num_P_tiles; t++) { + auto cO = rO[t].format(); + #pragma unroll + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp[r + (t % num_P_tiles)*REG_M]); } } + } else { + #pragma unroll + for(int t = 0; t < padded_head_size/REG_N*num_P_tiles; t++) { + auto cO = rO[t].format(); + #pragma unroll + for(int r = 0; r < REG_M; r++) + cO.row(r) = cm_mul(cO.row(r), max_comp[r + (t % num_P_tiles)*REG_M]); + } } - else { - //ugemm_PV1(slm_V, P, max_comp, rO, slm_offset); + + // ---- transpose each sub-tile and accumulate P@V into rO ---- + constexpr int num_Vchunks = padded_head_size/REG_N; + #pragma unroll + for(int b = 0; b < KV_BLK; b++) { + matrix P; + transpose_St_to_P_half(St.select(b*kv_step, 0), P); auto P2 = P.format(); + b2dV.set_block_y(kv_base + b*kv_step); #pragma unroll - for(int k = 0, ri=0; k < padded_head_size; k += REG_N, ri += num_P_tiles) { + for(int ci = 0, ri = 0; ci < num_Vchunks; ci++, ri += num_P_tiles) { matrix Vmat; - - cm_prefetch(prefetch_V.set_block_x(k)); - cm_load(Vmat.format(), b2dV.set_block_x(k)); - - //# compensate cur_O - // matrix rO; - #pragma unroll - for(int p = 0; p < num_P_tiles; p++) { - auto cO = rO[ri + p].format(); - #pragma unroll - for(int r = 0; r < REG_M; r++) - cO.row(r) = cm_mul(cO.row(r), max_comp[r + p*REG_M]); - } - + cm_load(Vmat.format(), b2dV.set_block_x(ci*REG_N)); #pragma unroll for(int p = 0; p < num_P_tiles; p++) { rO[ri + p] = cm_dpas( @@ -626,14 +399,15 @@ void sdpa_kernel( SurfaceIndex key [[type("buffer_t")]], SurfaceIndex value [[type("buffer_t")]], SurfaceIndex output [[type("buffer_t")]], + uint q_pitch_bytes, + uint k_pitch_bytes, + uint v_pitch_bytes, uint q_off, uint k_off, uint v_off, uint o_off) { constexpr int padded_head_size = (head_size + 16 - 1) / 16 *16; constexpr uint o_pitch = (num_heads * head_size * sizeof(half)); - constexpr uint q_pitch = is_qkv_fused ? ((num_heads + num_kv_heads*2) * head_size * sizeof(half)) : o_pitch; - constexpr uint kv_pitch = is_qkv_fused ? q_pitch : (num_kv_heads * head_size * sizeof(half)); vector cur_max; vector cur_sum; @@ -654,11 +428,11 @@ void sdpa_kernel( // load as many as possible given one address if constexpr (head_size == 128 || head_size == 64) { matrix QmatI32; - cm_load_2d(QmatI32, query, q_off, q_pitch); + cm_load_2d(QmatI32, query, q_off, q_pitch_bytes); #pragma unroll for(int k = 0, ri = 0; k < head_size/2; k += REG_K/2, ri++) { Transpose2DMatrix(QmatI32.select(0, k), rQ[ri].format()); - rQ[ri].format() = cm_mul(rQ[ri].format(), (half)scale_factor); + rQ[ri].format() = cm_mul(rQ[ri].format(), (half)q_prescale); } } else { constexpr int num_full_blocks = head_size / REG_K; @@ -667,9 +441,9 @@ void sdpa_kernel( for(; i < num_full_blocks; i++) { matrix QmatI32; int k = i * REG_K; - cm_load_2d(QmatI32, query, q_off + k * sizeof(uint) / 2, q_pitch); + cm_load_2d(QmatI32, query, q_off + k * sizeof(uint) / 2, q_pitch_bytes); Transpose2DMatrix(QmatI32, rQ[i].format()); - rQ[i].format() = cm_mul(rQ[i].format(), (half)scale_factor); + rQ[i].format() = cm_mul(rQ[i].format(), (half)q_prescale); } // if with tail, load with head_size_tail @@ -677,15 +451,16 @@ void sdpa_kernel( if constexpr (head_size % REG_K > 0) { int k = num_full_blocks * REG_K; matrix QmatI32; - cm_load_2d_with_tail(QmatI32, query, q_off + k * sizeof(half), q_pitch); + cm_load_2d_with_tail(QmatI32, query, q_off + k * sizeof(half), q_pitch_bytes); Transpose2DMatrix(QmatI32, rQ[num_full_blocks].format()); - rQ[num_full_blocks].format() = cm_mul(rQ[num_full_blocks].format(), (half)scale_factor); + rQ[num_full_blocks].format() = cm_mul(rQ[num_full_blocks].format(), (half)q_prescale); } } } constexpr int num_P_tiles = REG_N / REG_M; matrix rO; + rO = 0.0f; // Zero the accumulator: the first softmax block scales rO by max_comp==0, and 0*NaN==NaN if the GRF holds stale NaN/Inf bits. int causal_left = q_start; constexpr uint slm_buff_size = kv_step * padded_head_size * sizeof(half); @@ -693,7 +468,6 @@ void sdpa_kernel( int slm_buff_id_read = 0; auto load_slm_KV = [&](int kv_pos) { - //if (kv_pos < 1024000) return; int kv_tokens = kv_stop - kv_pos; if (kv_tokens <= 0) return; // Calculate valid rows for this block (used to zero out garbage data) @@ -703,14 +477,13 @@ void sdpa_kernel( // non-tail branch is faster if (wg_local_id < local_size/2) { - //if (kv_pos > 1024000) { matrix temp; // head_size is split into blocks of constexpr int num_full_blocks = head_size / REG_K; int i = wg_local_id; for (; i < num_full_blocks; i += local_size / 2) { int k = i * REG_K; - cm_load_2d(temp, key, k_off + k * sizeof(half), kv_pitch); + cm_load_2d(temp, key, k_off + k * sizeof(half), k_pitch_bytes); // Zero out unused K rows to prevent NaN from garbage data in KV cache // (Similar approach as PA kernel: cm_pa_common.hpp) for (int r = kv_valid_rows; r < kv_step; r++) @@ -723,7 +496,7 @@ void sdpa_kernel( // following code will be optimized out when head_size_tail = 0 if constexpr (head_size % REG_K > 0) { int k = num_full_blocks * REG_K; - cm_load_2d_with_tail<2*REG_M, REG_K, head_size % REG_K>(temp, key, k_off + k * sizeof(half), kv_pitch); + cm_load_2d_with_tail<2*REG_M, REG_K, head_size % REG_K>(temp, key, k_off + k * sizeof(half), k_pitch_bytes); // Zero out unused K rows for (int r = kv_valid_rows; r < kv_step; r++) temp.row(r) = 0; @@ -732,20 +505,18 @@ void sdpa_kernel( temp.format()); } } else { - //if (kv_pos > 1024000) { // read 16x16 XMX-B matrix (1x REG_N in Xe2, 2x REG_N in Xe1) constexpr int VK_STEP = 16; static_assert((VK_STEP % REG_N) == 0); matrix temp2; matrix temp_vnni; - //b2dV.set_block_y(kv_pos); // head_size is split into blocks of constexpr int num_full_blocks = head_size / VK_STEP; int i = wg_local_id - local_size / 2; for (; i < num_full_blocks; i += local_size / 2) { int k = i * VK_STEP; - cm_load_2d(temp2, value, v_off + k * sizeof(half), kv_pitch); + cm_load_2d(temp2, value, v_off + k * sizeof(half), v_pitch_bytes); // Zero out unused V rows to prevent NaN from garbage data in KV cache // (Similar approach as PA kernel: cm_pa_common.hpp) for (int r = kv_valid_rows; r < kv_step; r++) @@ -762,7 +533,7 @@ void sdpa_kernel( // following code will be optimized out when head_size_tail = 0 if constexpr (head_size % VK_STEP > 0) { int k = num_full_blocks * VK_STEP; - cm_load_2d_with_tail(temp2, value, v_off + k * sizeof(half), kv_pitch); + cm_load_2d_with_tail(temp2, value, v_off + k * sizeof(half), v_pitch_bytes); // Zero out unused V rows for (int r = kv_valid_rows; r < kv_step; r++) temp2.row(r) = 0; @@ -775,9 +546,8 @@ void sdpa_kernel( } } } - k_off += kv_step * kv_pitch; - v_off += kv_step * kv_pitch; - // printf(" diff= %lu\n", get_clock() - clk0); + k_off += kv_step * k_pitch_bytes; + v_off += kv_step * v_pitch_bytes; }; load_slm_KV(0); @@ -808,10 +578,8 @@ void sdpa_kernel( if (kv_pos + kv_step < kv_stop) cm_sbarrier(1); - //if (kv_pos < 1024000) continue; uint slm_offset = (slm_buff_id_read & 3) * slm_buff_size; - //=========================================================== 1807 ~ 3247 //# St = k @ Qt matrix St = ugemm_KQ(slm_K, rQ, slm_offset); @@ -829,7 +597,6 @@ void sdpa_kernel( St[p] = cm_add(St[p], cmask); if (v < q_step - 1) v++; } - //if (wg_local_id == 0) show(St);return; } causal_left -= kv_step; } @@ -838,8 +605,7 @@ void sdpa_kernel( int kv_tokens = kv_stop - kv_pos; for(int p = kv_tokens; p < kv_step; p++) St[p] = -3.4e38f; - //show(St); - auto max_comp = online_softmax_update(St, cur_max, cur_sum); + auto max_comp = cm_online_softmax_update(St, cur_max, cur_sum); matrix P; Transpose2DMatrix(St, P); @@ -865,7 +631,6 @@ void sdpa_kernel( cur_O_f16[r + p*REG_M] = cm_mul(cO.row(r), cur_sum[r + p*REG_M]); } } - // if (i == args_verbose) show(cur_O_f16); cm_store_2d(cur_O_f16, output, o_off + k*sizeof(half), o_pitch); } } diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp index 4d615fc1191b8c..27778bab936e04 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/paged_attention_gen.cpp @@ -381,6 +381,9 @@ JitConstants PagedAttentionGeneratorMultiToken::get_jit_constants(const kernel_i jit.make("CMPA_BLOCK_SZ", PA_KV_CACHE_BLOCK_SIZE_LEGACY); } jit.make("CMPA_WG_SEQ_LEN", get_wg_seq_len(params)); + // Tree softmax: wins on INT8 (dequant amortises scratch, shorter dep-chain helps), + // regresses ~9% on FP16 (extra scratch matrix increases register pressure). + jit.make("CMFLA_USE_TREE_SOFTMAX", get_kv_compressed(params) ? 1 : 0); return jit; } diff --git a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp index e4c12f182d0044..1b565836714b16 100644 --- a/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/cm/vl_sdpa_opt.cpp @@ -20,9 +20,39 @@ constexpr auto get_vlsdpa_build_options() { return " -cmc -Qxcm_register_file_size=256"; } -struct VLSDPARuntimeParams : public ImplRuntimeParams { - std::vector cu_seqlens; -}; +constexpr size_t round_up_to_tile(size_t value, size_t tile_size) { + return (value + tile_size - 1) / tile_size * tile_size; +} + +constexpr size_t get_default_kv_blk(size_t head_size) { + const size_t padded_head_size = round_up_to_tile(head_size, 16); + const size_t tail_size = padded_head_size - head_size; + + if (padded_head_size <= 64 || tail_size <= 1) { + return 2; + } + + return 1; +} + +// Enforce the CM kernel memory-layout contract documented in cm_sdpa_vlen.cm: +// Q/K/V/output layouts must all carry the fused transpose order {1,0,2}, i.e. the +// original SDPA-facing view [H, L, S] is expressed as physical memory [L, H, S] with +// the token (outer) axis at dim 0. When this contract is violated (empty or non-{1,0,2} +// orders, typically because TransposeVLSDPA fusion did not fire — e.g. surrounding +// Transposes were converted to Reshape by TransposeToReshape when a permuted dim +// degenerates to size 1), `get_pitches()[0]` no longer represents the per-token stride +// and the kernel silently walks out of bounds -> GPU hang (CL_OUT_OF_RESOURCES). +// Fail fast at compile / dispatch time with an actionable message. +inline void ensure_fused_transpose_order(const vl_sdpa& desc) { + const std::vector expected_order{1, 0, 2}; + OPENVINO_ASSERT(desc.input_q_transpose_order == expected_order && desc.input_k_transpose_order == expected_order && + desc.input_v_transpose_order == expected_order && desc.output_transpose_order == expected_order, + "VLSDPA CM impl requires the surrounding Transpose{1,0,2} to be fused into " + "VLSDPA (all four transpose orders must be {1,0,2}). If this fires, the " + "SDPAToVLSDPA / TransposeVLSDPAMatcher fusion did not run for this pattern; " + "kernel arg computation (`get_pitches()[0]` as the token stride) is unsafe."); +} class VLSDPAGenerator : public KernelGenerator { public: @@ -49,6 +79,7 @@ class VLSDPAGenerator : public KernelGenerator { }; auto desc = params.typed_desc(); + ensure_fused_transpose_order(*desc); const auto query_shape = transpose_pshape(params.get_input_layout(0).get_partial_shape(), desc->input_q_transpose_order); const auto key_shape = transpose_pshape(params.get_input_layout(1).get_partial_shape(), desc->input_k_transpose_order); @@ -56,6 +87,7 @@ class VLSDPAGenerator : public KernelGenerator { const size_t num_q_heads = query_shape[query_shape.size() - 3].get_length(); const size_t num_kv_heads = key_shape[key_shape.size() - 3].get_length(); const float scale_factor = 1.0f / std::sqrt(static_cast(head_size)); + const size_t cmfla_kv_blk = get_default_kv_blk(head_size); GPU_DEBUG_TRACE_DETAIL << "VLSDPA query_shape " << query_shape << ", q_transpose_order " << PartialShape(desc->input_q_transpose_order) << ", key_shape " << key_shape << ", k_transpose_order " << PartialShape(desc->input_k_transpose_order) @@ -67,6 +99,7 @@ class VLSDPAGenerator : public KernelGenerator { make_jit_constant("CMFLA_NUM_KV_HEADS", num_kv_heads), make_jit_constant("CMFLA_HEAD_SIZE", head_size), make_jit_constant("CMFLA_SCALE_FACTOR", scale_factor), + make_jit_constant("CMFLA_KV_BLK", cmfla_kv_blk), }); return jit; @@ -85,15 +118,26 @@ class VLSDPAGenerator : public KernelGenerator { args.push_back({ArgumentDescriptor::Types::INPUT, static_cast(params.input_layouts.size() - 1)}); // input: cu_seq_lens - args.push_back({ArgumentDescriptor::Types::SCALAR, 0}); + args.push_back({ArgumentDescriptor::Types::SCALAR, 0}); // need_wg_mapping + + // Per-input token-axis discontinuity descriptors (populated in the dispatch lambda + // from each input layout's get_linear_offset() and get_pitches()[0]). See + // cm_sdpa_vlen.cm for the exact memory-layout contract with the kernel. + args.push_back({ArgumentDescriptor::Types::SCALAR, 1}); // token_offset_q + args.push_back({ArgumentDescriptor::Types::SCALAR, 2}); // token_offset_k + args.push_back({ArgumentDescriptor::Types::SCALAR, 3}); // token_offset_v + args.push_back({ArgumentDescriptor::Types::SCALAR, 4}); // q_token_pitch + args.push_back({ArgumentDescriptor::Types::SCALAR, 5}); // k_token_pitch + args.push_back({ArgumentDescriptor::Types::SCALAR, 6}); // v_token_pitch return args; } [[nodiscard]] DispatchDataFunc get_dispatch_data_func() const override { - return DispatchDataFunc{[](const RuntimeParams& params, KernelData& kd, ImplRuntimeParams* rt_params) { + return DispatchDataFunc{[](const RuntimeParams& params, KernelData& kd, ImplRuntimeParams*) { assert(!params.is_dynamic()); auto desc = params.typed_desc(); + ensure_fused_transpose_order(*desc); // transpose shape into BHLS(4D), or HLS(3D) auto transpose_pshape = [](const ov::Shape& pshape, const std::vector& order) { @@ -107,28 +151,30 @@ class VLSDPAGenerator : public KernelGenerator { return transposed_pshape; }; const auto query_shape = transpose_pshape(params.get_input_layout(0).get_shape(), desc->input_q_transpose_order); - const size_t num_q_heads = query_shape[query_shape.size() - 3]; // TODO: make it to be configuration of primitive_inst + const size_t num_q_heads = query_shape[query_shape.size() - 3]; - const auto& vlsdpa_rt_params = static_cast(*rt_params); - const auto& cu_seqlens = vlsdpa_rt_params.cu_seqlens; + // Read cu_seqlens via mem_lock — zero-cost on USM-host (PC), no GPU sync stall. + const auto cu_seqlens_mem = params.memory_deps.at(params.input_layouts.size() - 1); + OPENVINO_ASSERT(cu_seqlens_mem->get_layout().data_type == ov::element::i32, "VLSDPA expects cu_seqlens input to be i32"); + mem_lock cu_seqlens_lock(cu_seqlens_mem, *params.strm); + const auto cu_seqlens_sz = cu_seqlens_lock.size(); + OPENVINO_ASSERT(cu_seqlens_sz >= 2, "VLSDPA expects cu_seqlens to contain at least two elements (start/end)"); size_t max_seq_len = 0; - for (size_t i = 1; i < cu_seqlens.size(); i++) { - auto start_idx = cu_seqlens[i - 1]; - auto end_idx = cu_seqlens[i]; + for (size_t i = 1; i < cu_seqlens_sz; i++) { + auto start_idx = cu_seqlens_lock[i - 1]; + auto end_idx = cu_seqlens_lock[i]; max_seq_len = std::max(max_seq_len, static_cast(end_idx - start_idx)); } const auto& info = params.get_device_info(); const size_t CM_GRF_WIDTH = (info.arch <= gpu_arch::xe_hpc) ? 256 : 512; - const size_t q_step = static_cast(std::floor(CM_GRF_WIDTH / 32)); // or 8 on Xe1 + const size_t q_step = static_cast(std::floor(CM_GRF_WIDTH / 32)); size_t wg_size = static_cast(std::floor((max_seq_len + q_step - 1) / q_step)); - int32_t need_wg_mapping = 0; + size_t need_wg_mapping = 0; if (wg_size > 16) { - // # seq_len is too big to fit into a single work-group - // # will use fixed work-group size 16, process 16*16 (or 16*8 on xe1) - // # part of sequence, in this case, kernel needs to figure-out which part - // # it needs to handle + // seq_len is too large for a single work-group; use wg_size=16 and let + // the kernel's while-loop scan cu_seqlens to find its sequence/block. need_wg_mapping = 1; wg_size = 16; } @@ -137,26 +183,45 @@ class VLSDPAGenerator : public KernelGenerator { if (need_wg_mapping) { wg_count = 0; const auto wg_seq_len = wg_size * q_step; - for (size_t i = 1; i < cu_seqlens.size(); i++) { - auto start_idx = cu_seqlens[i - 1]; - auto end_idx = cu_seqlens[i]; - wg_count += static_cast(std::floor((end_idx - start_idx + wg_seq_len - 1) / wg_seq_len)); + for (size_t i = 1; i < cu_seqlens_sz; i++) { + auto start_idx = cu_seqlens_lock[i - 1]; + auto end_idx = cu_seqlens_lock[i]; + wg_count += static_cast((end_idx - start_idx + wg_seq_len - 1) / wg_seq_len); } } else { - wg_count = cu_seqlens.size() - 1; + wg_count = cu_seqlens_sz - 1; } auto& wgs = kd.params.workGroups; wgs.global = {num_q_heads, wg_count * wg_size, 1}; wgs.local = {1, wg_size, 1}; - std::vector scalars{need_wg_mapping}; + // Element-count offsets and per-input token strides taken directly from the + // input layouts. Non-zero token_offset_x arises when an in-place crop + // (e.g. TransposeSplitMatcher axis=1) leaves the SVM base pointing at the + // packed parent buffer and expresses the slice view via layout padding on + // the token axis; the kernel shifts the Q/K/V pointers accordingly. + // Contract with the CM kernel (see cm_sdpa_vlen.cm): only the token (outer) + // axis of each input layout may be padded — the head and head_size strides + // are hardcoded in the kernel and are NOT parameterized here. + constexpr size_t i32_max = 0x7fffffff; + const auto token_offset_q = params.input_layouts[0].get_linear_offset(); + const auto token_offset_k = params.input_layouts[1].get_linear_offset(); + const auto token_offset_v = params.input_layouts[2].get_linear_offset(); + const size_t q_token_pitch = static_cast(params.input_layouts[0].get_pitches()[0]); + const size_t k_token_pitch = static_cast(params.input_layouts[1].get_pitches()[0]); + const size_t v_token_pitch = static_cast(params.input_layouts[2].get_pitches()[0]); + OPENVINO_ASSERT(token_offset_q <= i32_max && token_offset_k <= i32_max && token_offset_v <= i32_max && q_token_pitch <= i32_max && + k_token_pitch <= i32_max && v_token_pitch <= i32_max, + "VLSDPA token offsets/pitches must fit into int32"); + + std::vector scalars{need_wg_mapping, token_offset_q, token_offset_k, token_offset_v, q_token_pitch, k_token_pitch, v_token_pitch}; kd.params.scalars.clear(); for (auto i : scalars) { - scalar_desc desc; - desc.t = scalar_desc::Types::INT32; - desc.v.s32 = static_cast(i); - kd.params.scalars.push_back(desc); + scalar_desc s; + s.t = scalar_desc::Types::INT32; + s.v.s32 = static_cast(i); + kd.params.scalars.push_back(s); } }}; } @@ -168,9 +233,7 @@ class VLSDPAOptImpl : public PrimitiveImplCM { Stage::Ptr vl_sdpa = make_stage(); - VLSDPAOptImpl() : PrimitiveImplOCL(VLSDPAOptImplementationManager::get_type_info_static()) { - m_rt_params = std::make_unique(); - } + VLSDPAOptImpl() : PrimitiveImplOCL(VLSDPAOptImplementationManager::get_type_info_static()) {} VLSDPAOptImpl(const program_node& node, const RuntimeParams& params) : VLSDPAOptImpl() { add_stage(vl_sdpa, params); } @@ -178,27 +241,16 @@ class VLSDPAOptImpl : public PrimitiveImplCM { [[nodiscard]] std::unique_ptr clone() const override { return make_deep_copy(this); } - - void update_rt_params(const cldnn::primitive_inst& instance) override { - update_stages_flags(instance); - - auto rt_params = static_cast(m_rt_params.get()); - auto& vlsdpa_instance = dynamic_cast&>(instance); - vlsdpa_instance.get_mask_seqlens_from_memory(rt_params->cu_seqlens); - } - - cldnn::event::ptr execute(const std::vector& events, cldnn::primitive_inst& instance) override { - if (m_rt_params == nullptr) { - m_rt_params = std::make_unique(); - } - return PrimitiveImplCM::execute(events, instance); - } }; } // namespace std::unique_ptr VLSDPAOptImplementationManager::create_impl(const program_node& node, const RuntimeParams& params) const { assert(node.is_type()); + // Validate the fused-transpose invariant on the main thread, before triggering async + // kernel compilation — the same check inside get_jit_constants would be swallowed by + // the kernels cache and surface only as "Kernel for {vlsdpa:sdpa} is not found". + ensure_fused_transpose_order(*params.typed_desc()); return std::make_unique(node, params); } diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/paged_attention_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/paged_attention_opt.cpp index 5e41219aabe04f..0915210090772f 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/paged_attention_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/paged_attention_opt.cpp @@ -1799,8 +1799,8 @@ class PagedAttentionOptImpl : public SDPAImplBase { total_matrix_elements += evictable_size * evictable_size; total_vector_elements += evictable_size; } - GPU_DEBUG_TRACE_DETAIL << "Adaptive RKV: Allocating dynamic buffers - " - << "matrix: " << total_matrix_elements << ", vector: " << total_vector_elements << std::endl; + GPU_DEBUG_TRACE_DETAIL << "Adaptive RKV: Allocating dynamic buffers - " << "matrix: " << total_matrix_elements + << ", vector: " << total_vector_elements << std::endl; } else { // Fallback: use maximum size (512) if runtime values not available const size_t max_evictable_size = 512; diff --git a/src/plugins/intel_gpu/src/graph/include/vl_sdpa_inst.h b/src/plugins/intel_gpu/src/graph/include/vl_sdpa_inst.h index 42eaf052eec78b..4a86ade1283bef 100644 --- a/src/plugins/intel_gpu/src/graph/include/vl_sdpa_inst.h +++ b/src/plugins/intel_gpu/src/graph/include/vl_sdpa_inst.h @@ -18,7 +18,10 @@ struct typed_program_node : public typed_program_node_base { using parent::parent; program_node& input(size_t index = 0) const { return get_dependency(index); } - std::vector get_shape_infer_dependencies() const override { return {}; } + std::vector get_shape_infer_dependencies() const override { + // cu_seq_lens is input index 3; dispatch_data_func reads it via mem_lock. + return {3}; + } }; using vl_sdpa_node = typed_program_node; @@ -39,10 +42,6 @@ class typed_primitive_inst : public typed_primitive_inst_base static std::string to_string(const vl_sdpa_node& node); typed_primitive_inst(network& network, const vl_sdpa_node& node); - - void get_mask_seqlens_from_memory(std::vector& cu_seqlens) const; - - memory::ptr cu_seqlens_memory_ptr() const { return dep_memory_ptr(3); } }; using vl_sdpa_inst = typed_primitive_inst; diff --git a/src/plugins/intel_gpu/src/graph/vl_sdpa.cpp b/src/plugins/intel_gpu/src/graph/vl_sdpa.cpp index 60674bfe70bfe9..227a351264e3a2 100644 --- a/src/plugins/intel_gpu/src/graph/vl_sdpa.cpp +++ b/src/plugins/intel_gpu/src/graph/vl_sdpa.cpp @@ -5,7 +5,6 @@ #include "vl_sdpa_inst.h" #include "primitive_type_base.h" #include "json_object.h" -#include "intel_gpu/runtime/utils.hpp" namespace cldnn { GPU_DEFINE_PRIMITIVE_TYPE_ID(vl_sdpa); @@ -30,19 +29,4 @@ std::string vl_sdpa_inst::to_string(const vl_sdpa_node& node) { vl_sdpa_inst::typed_primitive_inst(network& network, const vl_sdpa_node& node) : parent(network, node) {} -void vl_sdpa_inst::get_mask_seqlens_from_memory(std::vector& cu_seqlens) const { - cldnn::stream& stream = get_network().get_stream(); - stream.finish(); - - const auto cu_seqlens_mem = cu_seqlens_memory_ptr(); - - auto size = cu_seqlens_mem->count(); - OPENVINO_ASSERT(cu_seqlens_mem->get_layout().data_type == ov::element::i32 && size > 0); - - cu_seqlens.resize(size); - cu_seqlens_mem->copy_to(stream, cu_seqlens.data(), 0, 0, size * sizeof(int32_t), true); - - GPU_DEBUG_TRACE_DETAIL << " get_mask_seqlens_from_memory " << cu_seqlens << std::endl; -} - } // namespace cldnn diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/vlsdpa_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/vlsdpa_gpu_test.cpp index c07786b1ab53f1..e45ca319ca8f66 100644 --- a/src/plugins/intel_gpu/tests/unit/test_cases/vlsdpa_gpu_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/test_cases/vlsdpa_gpu_test.cpp @@ -12,6 +12,8 @@ #include #include #include +#include +#include #include "scaled_dot_product_attention_inst.h" #include "vl_sdpa_inst.h" @@ -101,6 +103,69 @@ struct vlsdpa_gpu_test : public ::testing::TestWithParam { return topo; } + // Build a Split(axis=1, num_splits=3) crop chain fanning a fused [L, 3*H, S] + // parent into named "q"/"k"/"v" slices. Enabling optimize_data on the resulting + // topology causes prepare_buffer_fusing to keep the crops in-place, so each + // slice's runtime layout carries a non-zero linear_offset and the parent's row + // pitch (3 * num_heads * head_size). This is exactly the discontinuous-token-axis + // input shape that the CM vl_sdpa kernel's token_offset_{q,k,v} / {q,k,v}_token_pitch + // scalars are designed to handle (see cm_sdpa_vlen.cm memory-layout contract). + void add_qkv_split(topology& topo, + const primitive_id& fused_id, + const primitive_id& axis_data_id, + int32_t num_heads, + int32_t head_size) { + const auto slice_shape = tensor(batch(1), feature(num_heads), spatial(head_size, 1)); + topo.add(crop("q", { input_info(fused_id), input_info(axis_data_id) }, + slice_shape, tensor(feature(0 * num_heads)), + crop_ngraph_op_mode::split, 0, /*axis*/1, /*num_splits*/3)); + topo.add(crop("k", { input_info(fused_id), input_info(axis_data_id) }, + slice_shape, tensor(feature(1 * num_heads)), + crop_ngraph_op_mode::split, 1, /*axis*/1, /*num_splits*/3)); + topo.add(crop("v", { input_info(fused_id), input_info(axis_data_id) }, + slice_shape, tensor(feature(2 * num_heads)), + crop_ngraph_op_mode::split, 2, /*axis*/1, /*num_splits*/3)); + } + + topology create_ref_topology_fused(layout fused_layout, + layout attn_mask_layout, + memory::ptr axis_data_mem, + int32_t num_heads, + int32_t head_size) { + topology topo; + topo.add(input_layout("fused", fused_layout)); + topo.add(data("split_axis_ref", axis_data_mem)); + add_qkv_split(topo, "fused", "split_axis_ref", num_heads, head_size); + + topo.add(input_layout("attn", attn_mask_layout)); // "attention_mask" + + auto sdpa_prim = scaled_dot_product_attention("sdpa", {input_info("q"), input_info("k"), input_info("v"), input_info("attn")}, + false, -1, order_q, order_k, order_v, {0, 1, 2}, {}, false); + topo.add(sdpa_prim); + topo.add(permute("permute_o", input_info("sdpa"), order_o2)); + topo.add(reorder("result", input_info("permute_o"), format::bfyx, data_types::f16)); + return topo; + } + + topology create_topology_fused(layout fused_layout, + layout cu_seqlen_layout, + memory::ptr axis_data_mem, + int32_t num_heads, + int32_t head_size) { + topology topo; + topo.add(input_layout("fused", fused_layout)); + topo.add(data("split_axis_opt", axis_data_mem)); + add_qkv_split(topo, "fused", "split_axis_opt", num_heads, head_size); + + topo.add(input_layout("attn", cu_seqlen_layout)); // "cu_seqlens" + + auto vlsdpa_prim = vl_sdpa("vlsdpa", {input_info("q"), input_info("k"), input_info("v"), input_info("attn")}, + order_q, order_k, order_v, order_o); + topo.add(vlsdpa_prim); + topo.add(reorder("result", input_info("vlsdpa"), format::bfyx, data_types::f16)); + return topo; + } + std::tuple run_network(topology &topo, cldnn::memory::ptr q_mem, cldnn::memory::ptr k_mem, @@ -185,6 +250,78 @@ struct vlsdpa_gpu_test : public ::testing::TestWithParam { } } + void execute_fused_qkv(vlsdpa_test_params& p, const bool is_caching_test = false) { + const auto head_size = p.head_size; + const auto num_heads = p.num_heads; + const auto cu_seqlens = p.cu_seqlens; + const auto cumsum = cu_seqlens.back(); + const auto seq_length_q = cumsum; + const auto seq_length_kv = cumsum; + + assert(cu_seqlens.front() == 0); + + auto& engine = get_test_engine(); + + // Parent-layout dims: {L, 3*H, S}. After optimize_data-driven buffer fusing, + // the three Split(axis=1, num_splits=3) crops become in-place views of this + // parent, producing the discontinuous token-axis layout the CM kernel expects. + const int32_t fused_feature = 3 * num_heads; + cldnn::layout fused_dyn_layout({-1, fused_feature, head_size}, data_types::f16, format::bfyx); + cldnn::layout attn_mask_dyn_layout({1, -1, -1}, data_types::f16, format::bfyx); + cldnn::layout cu_seqlens_dyn_layout({-1}, data_types::i32, format::bfyx); + + // Split-axis constant shared by both topologies (Split expects axis as i64 scalar). + auto axis_data_mem = engine.allocate_memory({{}, data_types::i64, format::bfyx}); + set_values(axis_data_mem, {1}); + + topology ref_topo = create_ref_topology_fused(fused_dyn_layout, attn_mask_dyn_layout, + axis_data_mem, num_heads, head_size); + topology opt_topo = create_topology_fused(fused_dyn_layout, cu_seqlens_dyn_layout, + axis_data_mem, num_heads, head_size); + + // Concrete input allocations. The fused parent is filled once and reused by both + // networks, so any output divergence isolates the vl_sdpa kernel's handling of + // the non-zero token_offset_{q,k,v} / {q,k,v}_token_pitch scalar path. + cldnn::layout fused_static_layout({seq_length_q, fused_feature, head_size}, data_types::f16, format::bfyx); + cldnn::layout attn_mask_static_layout({1, seq_length_q, seq_length_kv}, data_types::f16, format::bfyx); + cldnn::layout cu_seqlens_static_layout({static_cast(cu_seqlens.size())}, data_types::i32, format::bfyx); + + auto fused_mem = engine.allocate_memory(fused_static_layout); + auto attn_mask_mem = engine.allocate_memory(attn_mask_static_layout); + auto cu_seqlens_mem = engine.allocate_memory(cu_seqlens_static_layout); + + load_input(fused_mem, 0); + get_attention_mask(attn_mask_mem, cu_seqlens); + get_cu_seqlens(cu_seqlens_mem, cu_seqlens); + + auto run = [&](topology& topo, cldnn::memory::ptr attn_mem, bool optimize) { + ExecutionConfig config = get_test_default_config(engine); + config.set_property(ov::intel_gpu::allow_new_shape_infer(true)); + // optimize=false keeps the reference crops materialized (contiguous Q/K/V). + // optimize=true lets prepare_buffer_fusing fold the vl_sdpa crops in-place, + // which is the whole point of this test. + config.set_property(ov::intel_gpu::optimize_data(optimize)); + + cldnn::network::ptr net = get_network(engine, topo, config, get_test_stream_ptr(), is_caching_test); + net->set_input_data("fused", fused_mem); + net->set_input_data("attn", attn_mem); + + auto outputs = net->execute(); + return outputs.at("result").get_memory(); + }; + + auto mem_ref_ptr = run(ref_topo, attn_mask_mem, /*optimize=*/false); + auto mem_opt_ptr = run(opt_topo, cu_seqlens_mem, /*optimize=*/true); + + cldnn::mem_lock ref_data(mem_ref_ptr, get_test_stream()); + cldnn::mem_lock opt_data(mem_opt_ptr, get_test_stream()); + for (size_t idx = 0; idx < ref_data.size(); idx++) { + ASSERT_FALSE(std::isnan(opt_data[idx]) || std::isnan(ref_data[idx])) << "NaN found at index " << idx; + } + auto ret = cosineSimilarity(ref_data, opt_data); + ASSERT_GE(ret, 0.95f); + } + static std::string PrintToStringParamName(const testing::TestParamInfo& info) { std::string result = "vlsdpa_gpu_test_" + std::to_string(info.param.head_size) + "_" + @@ -227,6 +364,25 @@ TEST_P(vlsdpa_gpu_test, basic_caching) { execute(p, true); } +// Exercises the token_offset_{q,k,v} / {q,k,v}_token_pitch runtime scalars that +// the CM vl_sdpa kernel gained to handle discontinuous Q/K/V slices of a fused +// [L, 3*H, S] parent buffer (Split(axis=1, num_splits=3) with in-place buffer fusing). +TEST_P(vlsdpa_gpu_test, discontinuous_qkv_fused) { + if (!check_vlsdpa_available()) + GTEST_SKIP(); + + auto p = GetParam(); + execute_fused_qkv(p); +} + +TEST_P(vlsdpa_gpu_test, discontinuous_qkv_fused_caching) { + if (!check_vlsdpa_available()) + GTEST_SKIP(); + + auto p = GetParam(); + execute_fused_qkv(p, true); +} + INSTANTIATE_TEST_SUITE_P(smoke_vlsdpa_gpu_test, vlsdpa_gpu_test, ::testing::Values(