From d974607299924cba91265a1536ac17648639c806 Mon Sep 17 00:00:00 2001 From: "Sabnis, Omkar" Date: Mon, 19 Oct 2026 10:18:58 -0700 Subject: [PATCH 1/8] [GPU] SDPA: use key-partitioned single-token kernels for up to 4 query rows --- .../src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp | 34 ++++++++++++++++--- .../src/graph/impls/ocl_v2/sdpa_opt.cl | 4 ++- .../tests/unit/test_cases/sdpa_gpu_test.cpp | 7 ++++ 3 files changed, 39 insertions(+), 6 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp index ba4f19418f8f..9f64b606dbe7 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp @@ -26,6 +26,30 @@ using namespace cldnn; namespace ov::intel_gpu::ocl { +namespace { + +// Decode with grouped-query heads stacked as query rows has a few rows per KV head (3-4 on common models). The single-token +// kernel splits the keys into partitions, so it stays parallel on long contexts, but it reads the keys once per row: +// from 8 rows on it was slower than the multi-token kernels. +constexpr int64_t max_single_token_rows = 4; + +int64_t target_rows(const RuntimeParams& params) { + auto desc = params.typed_desc(); + const auto target_seq_len = extract_dim(get_transposed_channel(ChannelName::Y, desc->input_q_transpose_order), params.input_layouts[0]); + return target_seq_len.is_static() ? target_seq_len.get_length() : -1; +} + +bool use_multi_token_stage(const RuntimeParams& params) { + return target_rows(params) > max_single_token_rows; +} + +// The micro single-token kernel has no key partitions: it is used for one query row only. +bool single_query_row(const RuntimeParams& params) { + return target_rows(params) <= 1; +} + +} // namespace + bool SDPAOpt::has_per_channel_compressed_kv(const kernel_impl_params& params) { const auto desc = params.typed_desc(); if (!desc->is_kv_compressed) { @@ -90,7 +114,7 @@ class SDPAOptImpl : public SDPAImplBase { } else { auto is_indirect = params.typed_desc()->indirect_axis != -1; GPU_DEBUG_TRACE_DETAIL << "add stage for non-dynamic, is_indirect = " << is_indirect << "\n"; - if (is_prefill_stage(params) || unaligned_head_size(params)) { + if (use_multi_token_stage(params) || unaligned_head_size(params)) { if (is_indirect) { GPU_DEBUG_TRACE_DETAIL << "add stage for indirect non-dynamic with prefill_stage \n"; add_stage(indirect_multi_tokens, params); @@ -115,7 +139,7 @@ class SDPAOptImpl : public SDPAImplBase { #ifdef ENABLE_ONEDNN_FOR_GPU const auto& gfx_ver = params.get_program().get_engine().get_device_info().gfx_ver; bool is_ARL_H = (gfx_ver.major == 12 && gfx_ver.minor == 74); - bool can_use_micro_sdpa = SDPAOpt::supports_micro_sdpa(params) && !is_ARL_H && !is_indirect; + bool can_use_micro_sdpa = SDPAOpt::supports_micro_sdpa(params) && !is_ARL_H && !is_indirect && single_query_row(params); if (can_use_micro_sdpa) { add_stage(regular_micro_single_token, params); } @@ -131,7 +155,7 @@ class SDPAOptImpl : public SDPAImplBase { [[nodiscard]] event::ptr execute(const std::vector& events, primitive_inst& instance) override { const auto& params = *instance.get_impl_params(); auto new_params = SDPABase::requires_shape_canonicalization(params) ? SDPABase::static_canonicalize_shapes(params) : params; - bool is_prefill = is_prefill_stage(new_params); + bool is_prefill = use_multi_token_stage(new_params); const bool use_multi_tokens = is_prefill || unaligned_head_size(new_params); bool is_indirect = need_indirect_load(static_cast(instance)); GPU_DEBUG_TRACE_DETAIL << "execute indirect = " << is_indirect << ", prefill = " << is_prefill << "\n"; @@ -160,7 +184,7 @@ class SDPAOptImpl : public SDPAImplBase { return execute_stage(events, instance, is_indirect ? indirect_multi_tokens : regular_multi_tokens); } #ifdef ENABLE_ONEDNN_FOR_GPU - if (has_stage(regular_micro_single_token) && !is_indirect) { + if (has_stage(regular_micro_single_token) && !is_indirect && single_query_row(new_params)) { return execute_stage(events, instance, regular_micro_single_token); } #endif @@ -182,7 +206,7 @@ class SDPAOptImpl : public SDPAImplBase { const auto head_size = get_head_size(params_canonicalization.get_input_layout(0), extended_input_q_transpose_order); const auto num_of_partitions = get_partitions_num(params_canonicalization, SDPAStage::SINGLE_TOKEN); - const auto is_prefill = is_prefill_stage(params_canonicalization); + const auto is_prefill = use_multi_token_stage(params_canonicalization); const size_t buf_elements_count = (num_of_partitions == 1 || is_prefill) ? 1 : params.output_layouts[0].count() / head_size * num_of_partitions; const size_t tmp_out_elements_count = (num_of_partitions == 1 || is_prefill) ? 2 : params.output_layouts[0].count() * num_of_partitions; diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl index 7321275b813d..4cdb8b4e1668 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl @@ -588,7 +588,9 @@ KERNEL(sdpa_opt)( #endif qk_val[seq_idx] += INPUT0_VAL_MIN; #elif !IS_CAUSAL && HAS_ATTN_MASK_INPUT - const uint attn_mask_offset = INPUT3_GET_INDEX_SAFE(b0_idx, b1_idx, target_seq_idx + seq_idx, start_partition_idx + seq_len); + // The _SAFE index macro does not parenthesize its arguments: without the parentheses the modulo + // applies to seq_idx only, and a mask broadcast over query rows is read out of bounds + const uint attn_mask_offset = INPUT3_GET_INDEX_SAFE(b0_idx, b1_idx, (target_seq_idx + seq_idx), (start_partition_idx + seq_len)); INPUT3_COMPUTE_TYPE mask_val = DECODE_INPUT3_COMPUTE_TYPE(attn_mask[attn_mask_offset]); #ifdef CLAMP_ATTN_MASK_INPUT // Conditionally clamp attention mask when attention mask differs from SOFTMAX_ACCUMULATOR_TYPE(f32) diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp index 45ed04a24c8f..41d13803e0dc 100644 --- a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp @@ -311,6 +311,13 @@ INSTANTIATE_TEST_SUITE_P( sdpa_test_params{64, 32, 128, 128, 2, true, 0.125f, false, 0.0f, data_types::bf16}, // bf16 scale_val sdpa_test_params{64, 32, 128, 128, 2, false, 1.0f, true, 0.5f}, // attn_mask only sdpa_test_params{512, 8, 1, 1024, 2, true}, + // a few query rows on a long context (e.g. stacked grouped-query heads at decode): single-token kernel, keys + // split into partitions; 4 rows is the limit, 5 takes the multi-token kernels + sdpa_test_params{64, 8, 4, 2048, 1, false}, + sdpa_test_params{64, 8, 4, 2048, 1, true}, + sdpa_test_params{64, 8, 3, 4096, 2, true}, + sdpa_test_params{128, 8, 4, 1024, 1, false}, + sdpa_test_params{128, 8, 5, 1024, 1, false}, sdpa_test_params{64, 32, 128, 128, 2, true, data_types::bf16}, // bf16 dynamic sdpa_test_params{64, 32, 128, 128, 2, false, data_types::bf16}, // bf16 static sdpa_test_params{64, 10, 77, 77, 1, true, data_types::bf16} // bf16 two ranks mask From 8d7b5d1187a1012c70ff9f9ea1ea3621d809f446 Mon Sep 17 00:00:00 2001 From: "Sabnis, Omkar" Date: Mon, 19 Oct 2026 17:12:44 -0700 Subject: [PATCH 2/8] [GPU] SDPA: optimize the code and improve comments --- .../src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp | 29 +++++++------------ .../graph/impls/ocl_v2/sdpa/sdpa_utils.hpp | 5 ++-- .../src/graph/impls/ocl_v2/sdpa_opt.cl | 3 +- .../tests/unit/test_cases/sdpa_gpu_test.cpp | 3 +- 4 files changed, 16 insertions(+), 24 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp index 067ace5d933f..889c837c78c0 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp @@ -28,24 +28,17 @@ namespace ov::intel_gpu::ocl { namespace { -// Decode with grouped-query heads stacked as query rows has a few rows per KV head (3-4 on common models). The single-token -// kernel splits the keys into partitions, so it stays parallel on long contexts, but it reads the keys once per row: -// from 8 rows on it was slower than the multi-token kernels. +// Decode with grouped-query heads stacked as query rows has 3-4 rows per KV head on common models. +// Single-token kernel reads the keys once per row (Arc B390, 4K keys: +0.3% at 6 rows, -9% at 8, -26% at 16). constexpr int64_t max_single_token_rows = 4; -int64_t target_rows(const RuntimeParams& params) { - auto desc = params.typed_desc(); - const auto target_seq_len = extract_dim(get_transposed_channel(ChannelName::Y, desc->input_q_transpose_order), params.input_layouts[0]); - return target_seq_len.is_static() ? target_seq_len.get_length() : -1; -} - bool use_multi_token_stage(const RuntimeParams& params) { - return target_rows(params) > max_single_token_rows; + return get_target_seq_len(params) > max_single_token_rows; } -// The micro single-token kernel has no key partitions: it is used for one query row only. +// Micro single-token kernel has no key partitions, so it is used for one query row only. bool single_query_row(const RuntimeParams& params) { - return target_rows(params) <= 1; + return get_target_seq_len(params) <= 1; } } // namespace @@ -155,10 +148,10 @@ class SDPAOptImpl : public SDPAImplBase { [[nodiscard]] event::ptr execute(const std::vector& events, primitive_inst& instance) override { const auto& params = *instance.get_impl_params(); auto new_params = SDPABase::requires_shape_canonicalization(params) ? SDPABase::static_canonicalize_shapes(params) : params; - bool is_prefill = use_multi_token_stage(new_params); - const bool use_multi_tokens = is_prefill || unaligned_head_size(new_params); + const bool multi_token_rows = use_multi_token_stage(new_params); + const bool use_multi_tokens = multi_token_rows || unaligned_head_size(new_params); bool is_indirect = need_indirect_load(static_cast(instance)); - GPU_DEBUG_TRACE_DETAIL << "execute indirect = " << is_indirect << ", prefill = " << is_prefill << "\n"; + GPU_DEBUG_TRACE_DETAIL << "execute indirect = " << is_indirect << ", multi_token_rows = " << multi_token_rows << "\n"; update_rt_params(instance); kernel_dump_info.clear_entries(); @@ -206,9 +199,9 @@ class SDPAOptImpl : public SDPAImplBase { const auto head_size = get_head_size(params_canonicalization.get_input_layout(0), extended_input_q_transpose_order); const auto num_of_partitions = get_partitions_num(params_canonicalization, SDPAStage::SINGLE_TOKEN); - const auto is_prefill = use_multi_token_stage(params_canonicalization); - const size_t buf_elements_count = (num_of_partitions == 1 || is_prefill) ? 1 : params.output_layouts[0].count() / head_size * num_of_partitions; - const size_t tmp_out_elements_count = (num_of_partitions == 1 || is_prefill) ? 2 : params.output_layouts[0].count() * num_of_partitions; + const auto multi_token_rows = use_multi_token_stage(params_canonicalization); + const size_t buf_elements_count = (num_of_partitions == 1 || multi_token_rows) ? 1 : params.output_layouts[0].count() / head_size * num_of_partitions; + const size_t tmp_out_elements_count = (num_of_partitions == 1 || multi_token_rows) ? 2 : params.output_layouts[0].count() * num_of_partitions; internal_buffers.emplace_back(buf_elements_count, ov::element::f32); internal_buffers.emplace_back(buf_elements_count, ov::element::f32); diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_utils.hpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_utils.hpp index f0b00c299bc5..9ac1139fb6f9 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_utils.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_utils.hpp @@ -205,10 +205,11 @@ inline ChannelName get_transposed_channel(ChannelName c, const std::vector(); const auto target_seq_len = extract_dim(get_transposed_channel(ChannelName::Y, desc->input_q_transpose_order), params.input_layouts[0]); - return target_seq_len.is_static() ? target_seq_len.get_length() > 1 : false; + return target_seq_len.is_static() ? target_seq_len.get_length() : -1; } inline bool unaligned_head_size(const RuntimeParams& params) { diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl index 4cdb8b4e1668..790132efb9f8 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl @@ -588,8 +588,7 @@ KERNEL(sdpa_opt)( #endif qk_val[seq_idx] += INPUT0_VAL_MIN; #elif !IS_CAUSAL && HAS_ATTN_MASK_INPUT - // The _SAFE index macro does not parenthesize its arguments: without the parentheses the modulo - // applies to seq_idx only, and a mask broadcast over query rows is read out of bounds + // _SAFE index macro doesn't parenthesize its arguments, so the modulo for a broadcast mask needs the whole sum const uint attn_mask_offset = INPUT3_GET_INDEX_SAFE(b0_idx, b1_idx, (target_seq_idx + seq_idx), (start_partition_idx + seq_len)); INPUT3_COMPUTE_TYPE mask_val = DECODE_INPUT3_COMPUTE_TYPE(attn_mask[attn_mask_offset]); #ifdef CLAMP_ATTN_MASK_INPUT diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp index 3dc18aca393d..04d0f0c6751b 100644 --- a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp @@ -313,8 +313,7 @@ INSTANTIATE_TEST_SUITE_P( sdpa_test_params{64, 32, 128, 128, 2, true, 0.125f, false, 0.0f, data_types::bf16}, // bf16 scale_val sdpa_test_params{64, 32, 128, 128, 2, false, 1.0f, true, 0.5f}, // attn_mask only sdpa_test_params{512, 8, 1, 1024, 2, true}, - // a few query rows on a long context (e.g. stacked grouped-query heads at decode): single-token kernel, keys - // split into partitions; 4 rows is the limit, 5 takes the multi-token kernels + // Few query rows (stacked GQA heads at decode): single-token kernel up to 4 rows, multi-token kernels from 5 sdpa_test_params{64, 8, 4, 2048, 1, false}, sdpa_test_params{64, 8, 4, 2048, 1, true}, sdpa_test_params{64, 8, 3, 4096, 2, true}, From 3db4909e0d7038e8f2d15b35a8c9f40429bf45a4 Mon Sep 17 00:00:00 2001 From: "Sabnis, Omkar" Date: Wed, 4 Nov 2026 14:41:54 -0800 Subject: [PATCH 3/8] [GPU] SDPA: keep non-exact few-row cases off the single-token kernel --- .../graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp | 1 + .../graph/impls/ocl_v2/sdpa/sdpa_gen_opt.hpp | 6 ++++++ .../src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp | 18 ++++++++++++----- .../src/graph/impls/ocl_v2/sdpa_opt.cl | 2 -- .../tests/unit/test_cases/sdpa_gpu_test.cpp | 20 ++++++++++++------- 5 files changed, 33 insertions(+), 14 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp index 6adf13c07042..794477f6f76b 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp @@ -300,6 +300,7 @@ JitConstants SDPAOptGeneratorFinalization::get_jit_constants(const kernel_impl_p auto jit = SDPAOptGeneratorBase::get_jit_constants_base(params, SDPAStage::FINALIZATION); jit.make("SDPA_STAGE_1", 1); jit.make("TARGET_SEQ_LEN_BLOCK_SIZE", get_target_seq_len_block_size()); + jit.make("MAX_PARTITIONS_NUM", get_max_partitions_num()); return jit; } diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.hpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.hpp index c1b8fcb512bd..f42c187807a2 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.hpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.hpp @@ -24,6 +24,12 @@ inline size_t get_target_seq_len_block_size() { return block_size; } +// The finalization stage keeps one value per key partition in local memory +inline size_t get_max_partitions_num() { + constexpr size_t max_partitions_num = 128; + return max_partitions_num; +} + inline size_t get_sg_number_scale_factor(const device_info& info, size_t head_size, size_t kernel_type) { constexpr size_t subgroup_size = 16; const size_t optimal_scale_factor = 2; diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp index 889c837c78c0..f4bdf93e72e1 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp @@ -32,8 +32,17 @@ namespace { // Single-token kernel reads the keys once per row (Arc B390, 4K keys: +0.3% at 6 rows, -9% at 8, -26% at 16). constexpr int64_t max_single_token_rows = 4; +// Several query rows stay on the multi-token kernels where the single-token stages are not exact for them: +// per-channel compressed KV needs the micro kernel, every key partition adds the sink to its softmax sum, +// and the finalization stage merges at most get_max_partitions_num() partitions. +bool single_token_supports_query_rows(const RuntimeParams& params) { + return !SDPAOpt::has_per_channel_compressed_kv(params) && !params.typed_desc()->has_sink_input && + get_partitions_num(params, SDPAStage::SINGLE_TOKEN) <= get_max_partitions_num(); +} + bool use_multi_token_stage(const RuntimeParams& params) { - return get_target_seq_len(params) > max_single_token_rows; + const auto query_rows = get_target_seq_len(params); + return query_rows > max_single_token_rows || (query_rows > 1 && !single_token_supports_query_rows(params)); } // Micro single-token kernel has no key partitions, so it is used for one query row only. @@ -193,14 +202,13 @@ class SDPAOptImpl : public SDPAImplBase { [[nodiscard]] std::vector get_internal_buffer_descs(const RuntimeParams& params) const override { std::vector internal_buffers; - auto desc = params.typed_desc(); - auto extended_input_q_transpose_order = extend_order_in_num_heads_dim(desc->input_q_transpose_order); auto params_canonicalization = SDPABase::requires_shape_canonicalization(params) ? SDPABase::static_canonicalize_shapes(params) : params; - const auto head_size = get_head_size(params_canonicalization.get_input_layout(0), extended_input_q_transpose_order); + // exp_sums and max_logits hold one value per output row and partition; the output is V_HEAD_SIZE wide + const auto v_head_size = SDPABase::get_effective_v_head_size(params_canonicalization); const auto num_of_partitions = get_partitions_num(params_canonicalization, SDPAStage::SINGLE_TOKEN); const auto multi_token_rows = use_multi_token_stage(params_canonicalization); - const size_t buf_elements_count = (num_of_partitions == 1 || multi_token_rows) ? 1 : params.output_layouts[0].count() / head_size * num_of_partitions; + const size_t buf_elements_count = (num_of_partitions == 1 || multi_token_rows) ? 1 : params.output_layouts[0].count() / v_head_size * num_of_partitions; const size_t tmp_out_elements_count = (num_of_partitions == 1 || multi_token_rows) ? 2 : params.output_layouts[0].count() * num_of_partitions; internal_buffers.emplace_back(buf_elements_count, ov::element::f32); diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl index 869e4511a433..5d62d9f013c4 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl @@ -2845,8 +2845,6 @@ KERNEL(sdpa_opt)( // max_logits [batch, heads_num, q_len, partition_idx] // tmp_out [batch, heads_num, q_len, partition_idx, head_size] -#define MAX_PARTITIONS_NUM 128 - REQD_SUB_GROUP_SIZE(SUBGROUP_SIZE) KERNEL(sdpa_opt_finalization_stage)( OPTIONAL_SHAPE_INFO_ARG diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp index be74ce13673b..987a6a2ec8ce 100644 --- a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp @@ -320,6 +320,7 @@ INSTANTIATE_TEST_SUITE_P( sdpa_test_params{64, 8, 3, 4096, 2, true}, sdpa_test_params{128, 8, 4, 1024, 1, false}, sdpa_test_params{128, 8, 5, 1024, 1, false}, + sdpa_test_params{64, 8, 3, 33024, 1, false}, // 129 key partitions: more than the finalization stage merges sdpa_test_params{64, 32, 128, 128, 2, true, data_types::bf16}, // bf16 dynamic sdpa_test_params{64, 32, 128, 128, 2, false, data_types::bf16}, // bf16 static sdpa_test_params{64, 10, 77, 77, 1, true, data_types::bf16} // bf16 two ranks mask @@ -494,21 +495,23 @@ TEST(sdpa_gpu_causal_mask, decode_32q_8kv_1024seq) { // against a host reference. Sink logits dominate the softmax denominator, so a missing renormalization // shows up as an output scale error that is checked element-wise (cosine similarity is scale-invariant). // The scalar mask keeps micro SDPA out, so the regular multi-token sdpa_opt kernel is exercised. -struct sdpa_gpu_sink_test : public ::testing::TestWithParam> { +// A few query rows over several key partitions must not take the partitioned single-token kernel, which adds the sink +// once per partition. +struct sdpa_gpu_sink_test : public ::testing::TestWithParam> { static std::string PrintToStringParamName(const testing::TestParamInfo& info) { - const auto [use_flashattn_v2, seq_kv] = info.param; - return std::string(use_flashattn_v2 ? "flashattn_v2" : "no_flashattn_v2") + "_kv" + std::to_string(seq_kv); + const auto [use_flashattn_v2, seq_q, seq_kv] = info.param; + return std::string(use_flashattn_v2 ? "flashattn_v2" : "no_flashattn_v2") + "_q" + std::to_string(seq_q) + "_kv" + + std::to_string(seq_kv); } }; TEST_P(sdpa_gpu_sink_test, matches_independent_reference) { - const auto [use_flashattn_v2, seq_kv] = GetParam(); + const auto [use_flashattn_v2, seq_q, seq_kv] = GetParam(); tests::random_generator rg; rg.set_seed(GET_SUITE_NAME); auto& engine = get_test_engine(); constexpr int num_heads = 4; - constexpr int seq_q = 32; constexpr int head_size = 64; const layout q_layout({1, num_heads, seq_q, head_size}, data_types::f16, format::bfyx); @@ -606,7 +609,7 @@ TEST_P(sdpa_gpu_sink_test, matches_independent_reference) { INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_sink, sdpa_gpu_sink_test, - ::testing::Combine(::testing::Bool(), ::testing::Values(32, 300)), + ::testing::Combine(::testing::Bool(), ::testing::Values(32, 3), ::testing::Values(32, 300)), sdpa_gpu_sink_test::PrintToStringParamName); struct micro_sdpa_prefetch_k_params { @@ -2003,7 +2006,10 @@ INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_compressed_kv_per_channel, sdpa_test_params{128, 40, 10, 512, 512, 1, 8, true}, sdpa_test_params{128, 40, 10, 1, 512, 1, 8, true}, sdpa_test_params{128, 40, 40, 512, 512, 1, 8, true}, - sdpa_test_params{128, 40, 40, 1, 512, 1, 8, true}), + sdpa_test_params{128, 40, 40, 1, 512, 1, 8, true}, + // Few query rows (stacked GQA heads at decode) need the micro kernel too + sdpa_test_params{128, 40, 10, 3, 512, 1, 4, false}, + sdpa_test_params{128, 40, 10, 3, 512, 1, 8, false}), sdpa_gpu_compressed_kv_per_channel_test::PrintToStringParamName); INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_compressed_kv_per_token, From 49d071b7bcad28085e12ffd316360937c725525f Mon Sep 17 00:00:00 2001 From: "Sabnis, Omkar" Date: Wed, 4 Nov 2026 14:42:12 -0800 Subject: [PATCH 4/8] [GPU] SDPA: handle key-broadcast attention masks without out-of-bounds reads --- src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl | 7 +++++-- .../intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp | 9 +++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl index 5d62d9f013c4..91041f6c72d3 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl @@ -1070,6 +1070,9 @@ inline MASK_VECTOR_TYPE FUNC(load_attn_mask)(OPTIONAL_SHAPE_INFO_ARG #endif #if !IS_CAUSAL && HAS_ATTN_MASK_INPUT const uint attn_mask_offset = INPUT3_GET_INDEX_SAFE(b0_idx, b1_idx, target_seq_idx, source_seq_idx); + // A mask broadcast over the keys (e.g. a scalar mask) has a single element along X; the _SAFE index above wraps + // to it, so the following keys must not step past it. + const uint mask_key_pitch = INPUT3_SIZE_X == 1 ? 0 : 1; if (target_seq_idx >= (uint)TARGET_SEQ_LEN) { unroll_for (uint i = 0; i < SUBGROUP_SIZE; i++) { mask_vec[i] = NAN; @@ -1077,7 +1080,7 @@ inline MASK_VECTOR_TYPE FUNC(load_attn_mask)(OPTIONAL_SHAPE_INFO_ARG } else { if (source_seq_idx + SUBGROUP_SIZE <= (uint)SOURCE_SEQ_LEN) { unroll_for (uint i = 0; i < SUBGROUP_SIZE; i++) { - const INPUT3_TYPE mask_val = attn_mask[attn_mask_offset + i]; + const INPUT3_TYPE mask_val = attn_mask[attn_mask_offset + i * mask_key_pitch]; mask_vec[i] = DECODE_INPUT3_COMPUTE_TYPE(mask_val); } } else { @@ -1088,7 +1091,7 @@ inline MASK_VECTOR_TYPE FUNC(load_attn_mask)(OPTIONAL_SHAPE_INFO_ARG const uint max_mask_offset = min(source_seq_idx + SUBGROUP_SIZE, (uint)SOURCE_SEQ_LEN); for (uint i = 0; i < SUBGROUP_SIZE; i++) { if (source_seq_idx + i < max_mask_offset) { - const INPUT3_TYPE mask_val = attn_mask[attn_mask_offset + i]; + const INPUT3_TYPE mask_val = attn_mask[attn_mask_offset + i * mask_key_pitch]; mask_vec[i] = DECODE_INPUT3_COMPUTE_TYPE(mask_val); } else { mask_vec[i] = NAN; diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp index 987a6a2ec8ce..ee67b2d3ebfe 100644 --- a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp @@ -528,13 +528,18 @@ TEST_P(sdpa_gpu_sink_test, matches_independent_reference) { auto q_mem = engine.allocate_memory(q_layout); auto k_mem = engine.allocate_memory(kv_layout); auto v_mem = engine.allocate_memory(kv_layout); - auto mask_mem = engine.allocate_memory(scalar_layout); + // The scalar mask is the first element of a buffer whose other elements are large: a kernel reading the + // broadcast mask past its single element adds them to the logits. + auto mask_storage = engine.allocate_memory(layout({64}, data_types::f16, format::bfyx)); + std::vector mask_storage_data(64, ov::float16(100.0f)); + mask_storage_data[0] = ov::float16(0.0f); + set_values(mask_storage, mask_storage_data); + auto mask_mem = engine.reinterpret_buffer(*mask_storage, scalar_layout); auto scale_mem = engine.allocate_memory(scalar_layout); auto sink_mem = engine.allocate_memory(sink_layout); set_values(q_mem, q_data); set_values(k_mem, k_data); set_values(v_mem, v_data); - set_values(mask_mem, {ov::float16(0.0f)}); set_values(scale_mem, {scale}); set_values(sink_mem, sink_data); From 5e0bf5acb14868d2fc43bcd9eb759611c565c71c Mon Sep 17 00:00:00 2001 From: "Sabnis, Omkar" Date: Mon, 9 Nov 2026 07:24:04 -0800 Subject: [PATCH 5/8] [GPU] SDPA: clamp the attention mask for every mask type A key partition masked only with -Inf has an infinite maximum, so its softmax becomes NaN, and finalization keeps the NaN when it scales that partition by zero. The clamp to the lowest finite value was skipped for f32 masks; apply it to every mask type. --- .../src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp | 8 -------- src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl | 4 +--- 2 files changed, 1 insertion(+), 11 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp index 794477f6f76b..381519a78871 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_gen_opt.cpp @@ -60,20 +60,12 @@ JitConstants SDPAOptGeneratorBase::get_jit_constants_base(const kernel_impl_para GPU_DEBUG_TRACE_DETAIL << "k_head_size = " << k_head_size << ", v_head_size = " << v_head_size << "\n"; size_t data_inputs_num = get_data_inputs_num(*desc); - size_t attn_mask_idx = ScaledDotProductAttentionInputIdx::ATTN_MASK; const bool has_attn_mask_input = sdpa_has_runtime_attn_mask_input(params); if (desc->attn_mask_val.has_value()) { jit.make("STATIC_SCALAR_ATTN_MASK_VALUE", desc->attn_mask_val.value()); jit.make("HAS_ATTN_MASK_INPUT", 0); } else { jit.make("HAS_ATTN_MASK_INPUT", has_attn_mask_input ? 1 : 0); - if (has_attn_mask_input) { - const auto& attn_mask_layout = params.get_input_layout(attn_mask_idx); - // Enable clamping only if attn_mask dtype differs from softmax_accumulator_type(f32) - if (attn_mask_layout.data_type != data_types::f32) { - jit.make("CLAMP_ATTN_MASK_INPUT", 1); - } - } } size_t scale_idx = ScaledDotProductAttentionInputIdx::SCALE; if ((data_inputs_num > scale_idx) && (!desc->scale_val.has_value())) { diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl index 91041f6c72d3..1e91f9c0f4f0 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl @@ -591,10 +591,8 @@ KERNEL(sdpa_opt)( // _SAFE index macro doesn't parenthesize its arguments, so the modulo for a broadcast mask needs the whole sum const uint attn_mask_offset = INPUT3_GET_INDEX_SAFE(b0_idx, b1_idx, (target_seq_idx + seq_idx), (start_partition_idx + seq_len)); INPUT3_COMPUTE_TYPE mask_val = DECODE_INPUT3_COMPUTE_TYPE(attn_mask[attn_mask_offset]); -#ifdef CLAMP_ATTN_MASK_INPUT - // Conditionally clamp attention mask when attention mask differs from SOFTMAX_ACCUMULATOR_TYPE(f32) + // A key partition masked only with -Inf would have an infinite maximum and turn its softmax into NaN mask_val = INPUT3_MAX_FUNC(mask_val, INPUT3_VAL_MIN); -#endif qk_val[seq_idx] += mask_val; #elif defined(STATIC_SCALAR_ATTN_MASK_VALUE) qk_val[seq_idx] += STATIC_SCALAR_ATTN_MASK_VALUE; From 055e02a375469cb6f83c0e6593aa6ed0cb4b070e Mon Sep 17 00:00:00 2001 From: "Sabnis, Omkar" Date: Mon, 9 Nov 2026 07:24:10 -0800 Subject: [PATCH 6/8] [GPU] SDPA: keep V wider than K on the multi-token kernels The single-token kernel keeps the query in local memory as wide as Q/K, but its query loads and its output reduction use the V width. With 2-4 query rows and V wider than K, stay on the multi-token kernels. --- .../intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp index f4bdf93e72e1..2c4c85af9560 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa/sdpa_opt.cpp @@ -34,10 +34,12 @@ constexpr int64_t max_single_token_rows = 4; // Several query rows stay on the multi-token kernels where the single-token stages are not exact for them: // per-channel compressed KV needs the micro kernel, every key partition adds the sink to its softmax sum, -// and the finalization stage merges at most get_max_partitions_num() partitions. +// the finalization stage merges at most get_max_partitions_num() partitions, and the query's local memory is K wide +// while its loads and the output reduction use the V width. bool single_token_supports_query_rows(const RuntimeParams& params) { return !SDPAOpt::has_per_channel_compressed_kv(params) && !params.typed_desc()->has_sink_input && - get_partitions_num(params, SDPAStage::SINGLE_TOKEN) <= get_max_partitions_num(); + get_partitions_num(params, SDPAStage::SINGLE_TOKEN) <= get_max_partitions_num() && + SDPABase::get_effective_v_head_size(params) <= SDPABase::get_effective_k_head_size(params); } bool use_multi_token_stage(const RuntimeParams& params) { From 5cfec253a4dd62931a3594933c939fb76b58ac01 Mon Sep 17 00:00:00 2001 From: "Sabnis, Omkar" Date: Mon, 9 Nov 2026 07:24:15 -0800 Subject: [PATCH 7/8] [GPU] SDPA: honor output padding in the finalization stage Finalization computed a flat output offset, so a padded output was written to the wrong elements when several key partitions are merged. Use SDPA_OUTPUT_GET_INDEX like the other output writes; it also covers the transposed output. --- .../intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl index 1e91f9c0f4f0..e74c6e3f1fdb 100644 --- a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl +++ b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/sdpa_opt.cl @@ -2933,15 +2933,7 @@ KERNEL(sdpa_opt_finalization_stage)( #endif acc += out_val_f * TO_SOFTMAX_ACCUMULATOR_TYPE(max_logits_u_exp_sum[partition_idx]); } - const uint out_offset = b0_idx * (NUM_HEADS * TARGET_SEQ_LEN * V_HEAD_SIZE) + -#ifdef OUTPUT_TRANSPOSE_ORDER_PRESENT - target_seq_idx * (NUM_HEADS * V_HEAD_SIZE) + - b1_idx * (V_HEAD_SIZE) + -#else - b1_idx * (TARGET_SEQ_LEN * V_HEAD_SIZE) + - target_seq_idx * (V_HEAD_SIZE) + -#endif - local_id; + const uint out_offset = SDPA_OUTPUT_GET_INDEX(b0_idx, b1_idx, target_seq_idx, local_id); output[out_offset] = TO_OUTPUT_TYPE(acc / global_exp_sum); } From b6a57b3d4de3346b349d025a117c51722e6688c8 Mon Sep 17 00:00:00 2001 From: "Sabnis, Omkar" Date: Mon, 9 Nov 2026 07:24:21 -0800 Subject: [PATCH 8/8] [GPU] SDPA: test few query rows against a host reference Cover the inputs the few-row route takes: fully masked key partitions and rows with f16 and f32 masks, causal masks, sliding window, partition edges, head sizes, asymmetric and unaligned widths, output transpose and padding, batch broadcast, bf16 and f32, and rows and keys changing between calls of one dynamic network. Each case also checks which kernel ran. Add 3-row cases to the by-token compressed KV tests. --- .../tests/unit/test_cases/sdpa_gpu_test.cpp | 301 +++++++++++++++++- 1 file changed, 298 insertions(+), 3 deletions(-) diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp index ee67b2d3ebfe..9ebbd93575df 100644 --- a/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp +++ b/src/plugins/intel_gpu/tests/unit/test_cases/sdpa_gpu_test.cpp @@ -617,6 +617,296 @@ INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_sink, ::testing::Combine(::testing::Bool(), ::testing::Values(32, 3), ::testing::Values(32, 300)), sdpa_gpu_sink_test::PrintToStringParamName); +// Inputs the few-row route opens up, against a host reference: masks, partition edges, widths, batches, output layouts. +enum class few_rows_mask { none, causal, causal_lower_right, causal_sliding, masked_tail, per_row_partition, whole_row, static_inf }; + +struct sdpa_few_rows_grid_params { + data_types dt; + data_types mask_dt; + int k_head_size; + int v_head_size; + int seq_q; + int seq_kv; + few_rows_mask mask; + int batch; + int kv_batch; + bool transposed_output; + bool padded_output; + bool dynamic; + const char* expected_kernel; +}; + +class sdpa_few_rows_grid_runner { +public: + static constexpr int num_heads = 2; + static constexpr int sliding_window = 8; + static constexpr const char* not_single_token = "not_single_token"; + + explicit sdpa_few_rows_grid_runner(const sdpa_few_rows_grid_params& p) : p(p), engine(get_test_engine()), rg("sdpa_few_rows_grid") { + const bool has_mask_input = p.mask == few_rows_mask::masked_tail || p.mask == few_rows_mask::per_row_partition || p.mask == few_rows_mask::whole_row; + const auto input = [&](int batch, int rows, int head) { + return layout({batch, num_heads, p.dynamic ? -1 : rows, head}, p.dt, format::bfyx); + }; + topology topo; + topo.add(input_layout("q", input(p.batch, p.seq_q, p.k_head_size))); + topo.add(input_layout("k", input(p.kv_batch, p.seq_kv, p.k_head_size))); + topo.add(input_layout("v", input(p.kv_batch, p.seq_kv, p.v_head_size))); + std::vector inputs = {input_info("q"), input_info("k"), input_info("v")}; + if (has_mask_input) { + topo.add(input_layout("mask", p.dynamic ? layout({1, 1, -1, -1}, p.mask_dt, format::bfyx) : layout({1, 1, p.seq_q, p.seq_kv}, p.mask_dt, format::bfyx))); + inputs.push_back(input_info("mask")); + } + const bool is_causal = p.mask == few_rows_mask::causal || p.mask == few_rows_mask::causal_lower_right || p.mask == few_rows_mask::causal_sliding; + const std::vector output_order = p.transposed_output ? std::vector{0, 2, 1, 3} : std::vector{0, 1, 2, 3}; + auto prim = scaled_dot_product_attention("sdpa", + inputs, + is_causal, + -1, + {0, 1, 2, 3}, + {0, 1, 2, 3}, + {0, 1, 2, 3}, + output_order, + {}, + false, + p.mask == few_rows_mask::causal_lower_right || p.mask == few_rows_mask::causal_sliding); + prim.scale_val = 1.0f; + if (p.mask == few_rows_mask::causal_sliding) { + prim.sliding_window = sliding_window; + } + if (p.mask == few_rows_mask::static_inf) { + prim.attn_mask_val = -std::numeric_limits::infinity(); + } + if (p.padded_output) { + prim.output_paddings = {padding({0, 0, 1, 0}, {0, 0, 1, 0})}; + } + topo.add(prim); + topo.add(reorder("result", input_info("sdpa"), format::bfyx, data_types::f32)); + + ExecutionConfig cfg = get_test_default_config(engine); + cfg.set_property(ov::intel_gpu::allow_new_shape_infer(true)); + cfg.set_property(ov::intel_gpu::force_implementations(ov::intel_gpu::ImplForcingMap{{"sdpa", {format::type::bfyx, "sdpa_opt"}}})); + net = get_network(engine, topo, cfg, get_test_stream_ptr(), false); + } + + // Runs one shape on the network and compares with the host reference; in a dynamic network the rows and keys can differ per call. + void run(int seq_q, int seq_kv, const char* expected_kernel) { + const auto q_data = fill("q", {p.batch, num_heads, seq_q, p.k_head_size}, p.dt); + const auto k_data = fill("k", {p.kv_batch, num_heads, seq_kv, p.k_head_size}, p.dt); + const auto v_data = fill("v", {p.kv_batch, num_heads, seq_kv, p.v_head_size}, p.dt); + + const auto is_masked = [&](int row, int key) { + switch (p.mask) { + case few_rows_mask::causal: + return key > row; + case few_rows_mask::causal_lower_right: + return key > std::max(0, seq_kv - seq_q) + row; + case few_rows_mask::causal_sliding: { + const int pos = std::max(0, seq_kv - seq_q) + row; + return key > pos || (pos >= sliding_window && key < pos - sliding_window + 1); + } + case few_rows_mask::masked_tail: + return key >= 256; + case few_rows_mask::per_row_partition: + return key / 256 == row; + case few_rows_mask::whole_row: + return row == 1; + case few_rows_mask::static_inf: + return true; + default: + return false; + } + }; + if (p.mask == few_rows_mask::masked_tail || p.mask == few_rows_mask::per_row_partition || p.mask == few_rows_mask::whole_row) { + auto mask_mem = engine.allocate_memory(layout({1, 1, seq_q, seq_kv}, p.mask_dt, format::bfyx)); + std::vector mask_data(static_cast(seq_q) * seq_kv, 0.0f); + for (int r = 0; r < seq_q; ++r) { + for (int j = 0; j < seq_kv; ++j) { + mask_data[static_cast(r) * seq_kv + j] = is_masked(r, j) ? -std::numeric_limits::infinity() : 0.0f; + } + } + if (p.mask_dt == data_types::f32) { + set_values(mask_mem, mask_data); + } else { + set_values(mask_mem, std::vector(mask_data.begin(), mask_data.end())); + } + net->set_input_data("mask", mask_mem); + } + + auto out_mem = net->execute().at("result").get_memory(); + + primitive_id sdpa_id; + for (const auto& info : net->get_primitives_info()) { + if (info.type_id == "scaled_dot_product_attention") { + sdpa_id = info.original_id; + } + } + auto sdpa_inst = net->get_primitive(sdpa_id); + const auto entries = sdpa_inst->get_impl()->get_kernels_dump_info(*sdpa_inst->get_impl_params()).get_entries(); + if (std::string(expected_kernel) == not_single_token) { + ASSERT_EQ(entries.find("_single_"), std::string::npos) << entries; + } else { + ASSERT_NE(entries.find(expected_kernel), std::string::npos) << entries; + } + + const size_t out_count = static_cast(p.batch) * num_heads * seq_q * p.v_head_size; + std::vector expected(out_count); + std::vector logits(seq_kv); + for (int b = 0; b < p.batch; ++b) { + const int kv_b = p.kv_batch == 1 ? 0 : b; + for (int h = 0; h < num_heads; ++h) { + for (int r = 0; r < seq_q; ++r) { + const auto* q = &q_data[((static_cast(b) * num_heads + h) * seq_q + r) * p.k_head_size]; + float max_logit = std::numeric_limits::lowest(); + bool any_unmasked = false; + for (int j = 0; j < seq_kv; ++j) { + any_unmasked = any_unmasked || !is_masked(r, j); + } + for (int j = 0; j < seq_kv; ++j) { + const auto* k = &k_data[((static_cast(kv_b) * num_heads + h) * seq_kv + j) * p.k_head_size]; + logits[j] = is_masked(r, j) ? std::numeric_limits::lowest() : std::inner_product(q, q + p.k_head_size, k, 0.0f); + max_logit = std::max(max_logit, logits[j]); + } + float denom = 0.0f; + for (auto& logit : logits) { + logit = logit == std::numeric_limits::lowest() ? 0.0f : std::exp(logit - max_logit); + denom += logit; + } + const size_t row_idx = p.transposed_output ? (static_cast(b) * seq_q + r) * num_heads + h : (static_cast(b) * num_heads + h) * seq_q + r; + auto* out = &expected[row_idx * p.v_head_size]; + if (!any_unmasked) { + std::fill_n(out, p.v_head_size, std::numeric_limits::quiet_NaN()); + continue; + } + for (int j = 0; j < seq_kv; ++j) { + const auto* v = &v_data[((static_cast(kv_b) * num_heads + h) * seq_kv + j) * p.v_head_size]; + for (int d = 0; d < p.v_head_size; ++d) { + out[d] += logits[j] / denom * v[d]; + } + } + } + } + } + + cldnn::mem_lock out_ptr(out_mem, get_test_stream()); + // A dynamic network can preallocate more rows than this call has. + ASSERT_GE(out_mem->get_layout().count(), expected.size()); + const float tolerance = p.dt == data_types::bf16 ? 3e-2f : 5e-3f; + for (size_t i = 0; i < expected.size(); ++i) { + // A row with every key masked has no defined result, it only has to stay finite. + if (std::isnan(expected[i])) { + ASSERT_TRUE(std::isfinite(out_ptr[i])) << "index " << i << " q " << seq_q << " kv " << seq_kv; + } else { + ASSERT_NEAR(expected[i], out_ptr[i], tolerance) << "index " << i << " q " << seq_q << " kv " << seq_kv; + } + } + } + + sdpa_few_rows_grid_params p; + cldnn::engine& engine; + tests::random_generator rg; + cldnn::network::ptr net; + +private: + std::vector fill(const primitive_id& id, const ov::PartialShape& shape, data_types dt) { + const layout l(shape, dt, format::bfyx); + auto mem = engine.allocate_memory(l); + auto data = rg.generate_random_1d(l.count(), -1.0f, 1.0f); + if (dt == data_types::f32) { + set_values(mem, data); + } else if (dt == data_types::bf16) { + std::vector stored(data.begin(), data.end()); + set_values(mem, stored); + data.assign(stored.begin(), stored.end()); + } else { + std::vector stored(data.begin(), data.end()); + set_values(mem, stored); + data.assign(stored.begin(), stored.end()); + } + net->set_input_data(id, mem); + return data; + } +}; + +struct sdpa_few_rows_grid_test : public ::testing::TestWithParam { + static std::string PrintToStringParamName(const testing::TestParamInfo& info) { + const auto& p = info.param; + return ov::element::Type(p.dt).get_type_name() + "_" + ov::element::Type(p.mask_dt).get_type_name() + "_k" + std::to_string(p.k_head_size) + "_v" + std::to_string(p.v_head_size) + "_q" + + std::to_string(p.seq_q) + "_kv" + std::to_string(p.seq_kv) + "_mask" + std::to_string(static_cast(p.mask)) + "_b" + std::to_string(p.batch) + + "x" + std::to_string(p.kv_batch) + (p.transposed_output ? "_outT" : "") + (p.padded_output ? "_pad" : "") + (p.dynamic ? "_dyn" : "_static"); + } +}; + +TEST_P(sdpa_few_rows_grid_test, matches_host_reference) { + const auto p = GetParam(); + sdpa_few_rows_grid_runner runner(p); + runner.run(p.seq_q, p.seq_kv, p.expected_kernel); +} + +static std::vector few_rows_grid() { + const auto f16 = data_types::f16; + const auto f32 = data_types::f32; + const auto multi = sdpa_few_rows_grid_runner::not_single_token; + std::vector cases; + const auto add = [&](data_types dt, data_types mask_dt, int k, int v, int q, int kv, few_rows_mask mask, int b = 1, int kv_b = 1, bool t = false, bool pad = false, const char* kernel = "_single_") { + cases.push_back({dt, mask_dt, k, v, q, kv, mask, b, kv_b, t, pad, false, kernel}); + cases.push_back({dt, mask_dt, k, v, q, kv, mask, b, kv_b, t, pad, true, kernel}); + }; + // Partition edges and head sizes + for (int kv : {17, 256, 257, 513}) { + add(f16, f16, 64, 64, 3, kv, few_rows_mask::none); + } + add(f16, f16, 128, 128, 3, 513, few_rows_mask::none); + add(f16, f16, 512, 512, 3, 1024, few_rows_mask::none); + // Fully masked key partitions and rows, -Inf with every mask type + for (const auto& [dt, mask_dt] : std::vector>{{f16, f16}, {f16, f32}, {f32, f32}}) { + add(dt, mask_dt, 64, 64, 3, 512, few_rows_mask::masked_tail); + add(dt, mask_dt, 64, 64, 3, 600, few_rows_mask::per_row_partition); + add(dt, mask_dt, 64, 64, 3, 600, few_rows_mask::whole_row); + } + add(f32, f32, 64, 64, 1, 512, few_rows_mask::masked_tail, 1, 1, false, false, ""); + add(f16, f16, 64, 64, 3, 600, few_rows_mask::causal); + add(f16, f16, 64, 64, 3, 600, few_rows_mask::causal_lower_right); + add(f16, f16, 64, 64, 3, 600, few_rows_mask::causal_sliding); + // Output layouts, batches, types + add(f16, f16, 64, 64, 3, 600, few_rows_mask::none, 2, 2, true); + add(f16, f16, 64, 64, 3, 600, few_rows_mask::none, 1, 1, false, true); + add(f16, f16, 64, 64, 3, 600, few_rows_mask::none, 2, 1); + add(data_types::bf16, f16, 64, 64, 3, 600, few_rows_mask::none); + // V narrower than K takes the route, wider or unaligned does not + add(f16, f16, 128, 64, 3, 600, few_rows_mask::none); + add(f16, f16, 64, 128, 3, 600, few_rows_mask::none, 1, 1, false, false, multi); + add(f16, f16, 64, 256, 3, 600, few_rows_mask::none, 1, 1, false, false, multi); + add(f16, f16, 40, 40, 3, 600, few_rows_mask::none, 1, 1, false, false, multi); + add(f16, f16, 64, 64, 5, 600, few_rows_mask::none, 1, 1, false, false, multi); + return cases; +} + +INSTANTIATE_TEST_SUITE_P(smoke_sdpa_few_rows_grid, + sdpa_few_rows_grid_test, + ::testing::ValuesIn(few_rows_grid()), + sdpa_few_rows_grid_test::PrintToStringParamName); + +// One dynamic network: internal buffers and stages must follow the rows and keys of each call. +TEST(sdpa_few_rows_grid_shapes, rows_and_keys_change_between_calls) { + sdpa_few_rows_grid_params p{data_types::f16, data_types::f16, 64, 64, 1, 512, few_rows_mask::none, 1, 1, false, false, true, "_single_"}; + sdpa_few_rows_grid_runner runner(p); + const std::vector> calls = {{1, 512, ""}, + {3, 512, "_single_"}, + {5, 512, sdpa_few_rows_grid_runner::not_single_token}, + {3, 1024, "_single_"}, + {4, 300, "_single_"}, + {1, 1024, ""}, + {2, 33024, sdpa_few_rows_grid_runner::not_single_token}, + {3, 512, "_single_"}}; + for (const auto& [q, kv, kernel] : calls) { + SCOPED_TRACE("q " + std::to_string(q) + " kv " + std::to_string(kv)); + runner.run(q, kv, kernel); + if (::testing::Test::HasFatalFailure()) { + return; + } + } +} + struct micro_sdpa_prefetch_k_params { int head_size; int num_heads; @@ -2033,7 +2323,10 @@ INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_compressed_kv_per_token, // f32 compute: dequantized f16 KV must be cast to the Q compute type sdpa_test_params{128, 40, 10, 512, 512, 1, 8, false, data_types::f32}, sdpa_test_params{128, 40, 10, 128, 128, 1, 8, true, data_types::f32}, - sdpa_test_params{128, 40, 40, 128, 128, 1, 8, true, data_types::f32}), + sdpa_test_params{128, 40, 40, 128, 128, 1, 8, true, data_types::f32}, + sdpa_test_params{128, 40, 10, 3, 512, 1, 8, false}, + sdpa_test_params{128, 40, 40, 3, 512, 1, 8, false}, + sdpa_test_params{128, 40, 10, 3, 600, 1, 8, true}), sdpa_gpu_compressed_kv_per_token_test::PrintToStringParamName); INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_compressed_kv_bf16, @@ -2046,7 +2339,8 @@ INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_compressed_kv_bf16, sdpa_test_params{128, 40, 10, 128, 128, 1, 8, true}, sdpa_test_params{128, 40, 10, 1, 512, 1, 8, true}, sdpa_test_params{128, 40, 40, 128, 128, 1, 8, true}, - sdpa_test_params{128, 40, 40, 1, 512, 1, 8, true}), + sdpa_test_params{128, 40, 40, 1, 512, 1, 8, true}, + sdpa_test_params{128, 40, 10, 3, 600, 1, 8, false}), sdpa_gpu_compressed_kv_bf16_test::PrintToStringParamName); INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_compressed_kv_int4_bf16_per_token, @@ -2054,7 +2348,8 @@ INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_compressed_kv_int4_bf16_per_token, ::testing::Values(sdpa_test_params{128, 40, 10, 128, 128, 1, 4, true}, sdpa_test_params{128, 40, 10, 1, 512, 1, 4, true}, sdpa_test_params{128, 40, 40, 128, 128, 1, 4, true}, - sdpa_test_params{128, 40, 40, 1, 512, 1, 4, true}), + sdpa_test_params{128, 40, 40, 1, 512, 1, 4, true}, + sdpa_test_params{128, 40, 10, 3, 600, 1, 4, true}), sdpa_gpu_compressed_kv_int4_bf16_per_token_test::PrintToStringParamName); INSTANTIATE_TEST_SUITE_P(smoke_sdpa_gpu_compressed_kv_bf16_int8_zp,