[GPU] Enable SDPA optimization by default (#24767)

### Details:
This PR is a copy of
https://github.com/openvinotoolkit/openvino/pull/24757 from release
branch, it includes:
- Enabled SDPA by default
- Updated SDPA decomposition rule to cover only well-checked cases
- Updated functional tests accordingly
This commit is contained in:
Sergey Shlyapnikov 2024-05-29 20:36:42 +04:00 committed by GitHub
parent 45e219b450
commit cfb5d4ee84
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 17 additions and 36 deletions

View File

@ -321,10 +321,6 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
const auto& value_ps = sdpa->get_input_partial_shape(2);
// Known limitations:
// - SDPA impl could be slower in non-LLM scenarios than decomposed version
if (func->get_variables().size() == 0)
return false;
// - The data type of SDPA should be fp16
if (sdpa->get_output_element_type(0) != ov::element::f16)
return false;
@ -347,7 +343,7 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
// - The head size should be divisible by 16
const auto optimal_subgroup_size = 16;
if (query_ps[query_ps.size() - 1].is_dynamic() ||
query_ps[query_ps.size() - 1].get_length() > 256 ||
query_ps[query_ps.size() - 1].get_length() != 128 ||
query_ps[query_ps.size() - 1].get_length() % optimal_subgroup_size != 0) {
return false;
}

View File

@ -50,7 +50,7 @@ void ExecutionConfig::set_default() {
std::make_tuple(ov::intel_gpu::hint::host_task_priority, ov::hint::Priority::MEDIUM),
std::make_tuple(ov::intel_gpu::hint::queue_throttle, ov::intel_gpu::hint::ThrottleLevel::MEDIUM),
std::make_tuple(ov::intel_gpu::hint::queue_priority, ov::hint::Priority::MEDIUM),
std::make_tuple(ov::intel_gpu::hint::enable_sdpa_optimization, false),
std::make_tuple(ov::intel_gpu::hint::enable_sdpa_optimization, true),
std::make_tuple(ov::intel_gpu::enable_loop_unrolling, true),
std::make_tuple(ov::intel_gpu::disable_winograd_convolution, false),
std::make_tuple(ov::internal::exclusive_async_requests, false),

View File

@ -106,21 +106,9 @@ void ScaledAttnLayerGPUTest::SetUp() {
}
}
// Add artificial read/value operations to the model to trigger the enabling of the SDPA operation
auto read_key = std::make_shared<ov::op::v3::ReadValue>(inputParams.at(1), "v0");
auto assign_key = std::make_shared<ov::op::v3::Assign>(read_key, "v0");
auto read_value = std::make_shared<ov::op::v3::ReadValue>(inputParams.at(2), "v0");
auto assign_value = std::make_shared<ov::op::v3::Assign>(read_value, "v0");
ov::OutputVector inputs;
for (size_t i = 0; i < inputParams.size(); i++) {
if (i == 1)
inputs.push_back(read_key);
else if (i == 2)
inputs.push_back(read_value);
else
inputs.push_back(inputParams[i]);
inputs.push_back(inputParams[i]);
}
auto sdp = std::make_shared<ov::opset13::ScaledDotProductAttention>(inputs, is_causal);
@ -128,7 +116,7 @@ void ScaledAttnLayerGPUTest::SetUp() {
auto output = std::make_shared<ov::op::v0::Result>(sdp->output(0));
function = std::make_shared<ov::Model>(ov::OutputVector{output}, ov::SinkVector{assign_key, assign_value}, inputParams, "sdpa_model");
function = std::make_shared<ov::Model>(ov::OutputVector{output}, inputParams, "sdpa_model");
functionRefs = function->clone();
ov::pass::Manager manager;
@ -137,11 +125,8 @@ void ScaledAttnLayerGPUTest::SetUp() {
manager.register_pass<ov::pass::ScaledDotProductAttentionDecomposition>();
manager.run_passes(functionRefs);
// Enable SDPA
configuration.insert(ov::intel_gpu::hint::enable_sdpa_optimization(true));
auto it = std::find_if(inputShapes[1].second.begin(), inputShapes[1].second.end(), [&](const ov::Shape& shape){
return shape[2] >= 384;
return shape[2] >= 384 || shape[3] >= 128;
});
bool has_long_seq = it != inputShapes[1].second.end();
@ -190,12 +175,12 @@ const std::vector<std::vector<InputShape>> shapes{
// normal case, shapes of q,k,v are same
{
// q shape
{ov::test::InputShape{ov::PartialShape{-1, 8, -1, 64},
{ov::Shape{1, 8, 100, 64}, ov::Shape{1, 8, 1, 64}, ov::Shape{2, 8, 10, 64}}}
{ov::test::InputShape{ov::PartialShape{-1, 8, -1, 128},
{ov::Shape{1, 8, 100, 128}, ov::Shape{1, 8, 1, 128}, ov::Shape{2, 8, 10, 128}}}
},
// kv shape
{ov::test::InputShape{ov::PartialShape{-1, 8, -1, 64},
{ov::Shape{1, 8, 100, 64}, ov::Shape{1, 8, 1, 64}, ov::Shape{2, 8, 10, 64}}}
{ov::test::InputShape{ov::PartialShape{-1, 8, -1, 128},
{ov::Shape{1, 8, 100, 128}, ov::Shape{1, 8, 1, 128}, ov::Shape{2, 8, 10, 128}}}
},
// attn shape: [B, 1, -1, L0+L1]
{ov::test::InputShape{ov::PartialShape{-1, 1, -1, -1},
@ -204,12 +189,12 @@ const std::vector<std::vector<InputShape>> shapes{
},
{
// q shape
{ov::test::InputShape{ov::PartialShape{-1, 5, -1, 64},
{ov::Shape{2, 5, 100, 64}, ov::Shape{2, 5, 1, 64}, ov::Shape{2, 5, 384, 64}}}
{ov::test::InputShape{ov::PartialShape{-1, 5, -1, 128},
{ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 384, 128}}}
},
// kv shape
{ov::test::InputShape{ov::PartialShape{-1, 5, -1, 64},
{ov::Shape{2, 5, 100, 64}, ov::Shape{2, 5, 1, 64}, ov::Shape{2, 5, 384, 64}}}
{ov::test::InputShape{ov::PartialShape{-1, 5, -1, 128},
{ov::Shape{2, 5, 100, 128}, ov::Shape{2, 5, 1, 128}, ov::Shape{2, 5, 384, 128}}}
},
// attn shape: [B, 1, -1, L0+L1]
{ov::test::InputShape{ov::PartialShape{-1, 1, -1, -1},
@ -219,12 +204,12 @@ const std::vector<std::vector<InputShape>> shapes{
// heads number of kv is 1, attn mask: [B, H, L1, L0+L1]
{
// q shape
{ov::test::InputShape{ov::PartialShape{-1, 8, -1, 64},
{ov::Shape{1, 8, 100, 64}, ov::Shape{1, 8, 1, 64}, ov::Shape{2, 8, 10, 64}}}
{ov::test::InputShape{ov::PartialShape{-1, 8, -1, 128},
{ov::Shape{1, 8, 100, 128}, ov::Shape{1, 8, 1, 128}, ov::Shape{2, 8, 10, 128}}}
},
// kv shape
{ov::test::InputShape{ov::PartialShape{-1, 1, -1, 64},
{ov::Shape{1, 1, 100, 64}, ov::Shape{1, 1, 1, 64}, ov::Shape{2, 1, 10, 64}}}
{ov::test::InputShape{ov::PartialShape{-1, 1, -1, 128},
{ov::Shape{1, 1, 100, 128}, ov::Shape{1, 1, 1, 128}, ov::Shape{2, 1, 10, 128}}}
},
// attn shape
{ov::test::InputShape{ov::PartialShape{-1, 8, -1, -1},