[GPU] Enable int8/int4 quantized-KV GroupQueryAttention (openvinotoolkit#36798)

ankitm3k · web-flow · commit ea600b76705c · 2026-07-15T13:57:16.000Z
### Details: Enables the com.microsoft GroupQueryAttention op with an i8/u8 quantized KV cache on the Intel GPU plugin, for both the dynamic (Slice+Concat) and static (ScatterUpdate) cache-assembly paths. This already worked on CPU/TEMPLATE via the device-agnostic GroupQueryAttentionDecomposition; the GPU failures were a few plugin-local primitive gaps where a kernel supported float but not i8/u8. Fixes (all GPU-plugin-local; core op and decomposition untouched): - **KeepGQAKVScalePrecision**: new GPU MatcherPass that marks the KV dequant scales (k_scale/v_scale) precision-sensitive before ConvertPrecision. GPU runs ConvertPrecision (fp32->fp16) before the GQA op is decomposed in CommonOptimizations, and the intact op requires fp32 scales (com.microsoft spec). Mirrors the existing KeepXAttentionThresholdPrecision. Marks only the scale constants (no downstream fp32 propagation), so no extra memory or compute. - **ScatterUpdate (static path)**: adds i8/u8 to the OCL impl's supported input types and to the ref kernel's GetSupportedKey inputs (outputs already had them). Pure index-driven copy, so i8/u8 is numerically safe. - **Slice (dynamic i8 path)**: adds i8 to the ref kernel's supported output types (i8/u8 input and u8 output were already present). Tests: un-skips the 4 IE_GPU gqa i8/i4 conformance tests and adds static-shape (ScatterUpdate, reshape-to-static) variants; present KV is checked bit-exact. Run: `ov_onnx_frontend_tests --gtest_filter='*onnx_model_gqa_i8kv*:*onnx_model_gqa_i4kv*'` → 27/27 (IE_CPU 9, IE_GPU 9, INTERPRETER 9), no CPU/TEMPLATE regression. Out of scope: f8e4m3 KV on GPU (needs an f8e4m3 Slice/reorder layout) — its 2 IE_GPU tests remain skipped as a documented follow-up. ### Tickets: - [CVS-190527](https://jira.devtools.intel.com/browse/CVS-190527) ### AI Assistance: - *AI assistance used: no / yes* - *If yes, summarize how AI was used and what human validation was performed (build/tests/manual checks).*
diff --git a/src/frontends/onnx/tests/onnx_import_com_microsoft.in.cpp b/src/frontends/onnx/tests/onnx_import_com_microsoft.in.cpp
diff --git a/src/frontends/onnx/tests/runtime/ie/unit_test.manifest b/src/frontends/onnx/tests/runtime/ie/unit_test.manifest
@@ -359,13 +359,10 @@ IE_GPU.onnx_model_gather_elements_int32_axis_0
 IE_GPU.onnx_model_gather_elements_int8_axis_1
 IE_GPU.onnx_model_gather_elements_float_3D_axis_2
 
-# GroupQueryAttention with a quantized KV cache appends tokens via ScatterUpdate on i8/u8/f8e4m3
-# data, for which the GPU plugin has no layout ("No layout format available for scatterupdate
-# ... data_type: i8"). The numerics are covered on CPU/INTERPRETER; GPU is skipped here.
-IE_GPU.onnx_model_gqa_i8kv_per_channel
-IE_GPU.onnx_model_gqa_i8kv_per_tensor
-IE_GPU.onnx_model_gqa_i4kv_per_channel
-IE_GPU.onnx_model_gqa_i4kv_per_tensor
+# GroupQueryAttention with an f8e4m3 quantized KV cache: the GPU plugin has no f8e4m3 Slice/reorder
+# layout yet ("No layout format available"), so the fp8 variants remain skipped on GPU. The i8/i4
+# variants are enabled (see KeepGQAKVScalePrecision + i8/u8 ScatterUpdate/Slice support). fp8 GPU is
+# tracked as a separate follow-up; numerics stay covered on CPU/INTERPRETER meanwhile.
 IE_GPU.onnx_model_gqa_f8e4m3fnkv_per_channel
 IE_GPU.onnx_model_gqa_f8e4m3fnkv_per_tensor
 
diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl/scatter_update.hpp b/src/plugins/intel_gpu/src/graph/impls/ocl/scatter_update.hpp
@@ -37,11 +37,11 @@ struct ScatterUpdateImplementationManager : public ImplementationManager {
             format::bfwzyx
         };
 
-        static const std::vector<ov::element::Type_t> supported_in_types = {
-            ov::element::f32,
-            ov::element::f16,
-            ov::element::i32
-        };
+        static const std::vector<ov::element::Type_t> supported_in_types = {ov::element::f32,
+                                                                            ov::element::f16,
+                                                                            ov::element::i32,
+                                                                            ov::element::i8,
+                                                                            ov::element::u8};
 
         static const std::vector<ov::element::Type_t> supported_out_types = {
             ov::element::f32,
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/scatter_update/scatter_update_kernel_ref.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/scatter_update/scatter_update_kernel_ref.cpp
@@ -37,6 +37,8 @@ ParamsKey ScatterUpdateKernelRef::GetSupportedKey() const {
     k.EnableInputDataType(Datatype::F16);
     k.EnableInputDataType(Datatype::F32);
     k.EnableInputDataType(Datatype::INT32);
+    k.EnableInputDataType(Datatype::INT8);
+    k.EnableInputDataType(Datatype::UINT8);
     k.EnableOutputDataType(Datatype::F16);
     k.EnableOutputDataType(Datatype::F32);
     k.EnableOutputDataType(Datatype::INT32);
diff --git a/src/plugins/intel_gpu/src/kernel_selector/kernels/slice/slice_kernel_ref.cpp b/src/plugins/intel_gpu/src/kernel_selector/kernels/slice/slice_kernel_ref.cpp
@@ -90,6 +90,7 @@ ParamsKey SliceKernelRef::GetSupportedKey() const {
     k.EnableOutputDataType(Datatype::F32);
     k.EnableOutputDataType(Datatype::INT32);
     k.EnableOutputDataType(Datatype::INT64);
+    k.EnableOutputDataType(Datatype::INT8);
     k.EnableOutputDataType(Datatype::UINT8);
     k.EnableInputLayout(DataLayout::bfyx);
     k.EnableInputLayout(DataLayout::bfzyx);
diff --git a/src/plugins/intel_gpu/src/plugin/transformations/keep_gqa_kv_scale_precision.cpp b/src/plugins/intel_gpu/src/plugin/transformations/keep_gqa_kv_scale_precision.cpp
@@ -0,0 +1,39 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "keep_gqa_kv_scale_precision.hpp"
+
+#include "openvino/op/group_query_attention.hpp"
+#include "openvino/op/util/precision_sensitive_attribute.hpp"
+#include "openvino/pass/pattern/op/wrap_type.hpp"
+
+namespace ov::intel_gpu {
+
+KeepGQAKVScalePrecision::KeepGQAKVScalePrecision() {
+    auto gqa_m = ov::pass::pattern::wrap_type<ov::op::internal::GroupQueryAttention>();
+
+    ov::matcher_pass_callback callback = [=](ov::pass::pattern::Matcher& m) {
+        auto gqa = ov::as_type_ptr<ov::op::internal::GroupQueryAttention>(m.get_match_root());
+        if (!gqa)
+            return false;
+
+        if (transformation_callback(gqa))
+            return false;
+
+        // Only a quantized KV cache has scales at inputs 12 (k_scale) / 13 (v_scale).
+        if (!gqa->is_kv_quantized())
+            return false;
+
+        constexpr size_t k_scale_idx = 12;
+        constexpr size_t v_scale_idx = 13;
+        ov::mark_as_precision_sensitive(gqa->input(k_scale_idx));
+        ov::mark_as_precision_sensitive(gqa->input(v_scale_idx));
+        return true;
+    };
+
+    auto matcher = std::make_shared<ov::pass::pattern::Matcher>(gqa_m, "KeepGQAKVScalePrecision");
+    register_matcher(matcher, callback);
+}
+
+}  // namespace ov::intel_gpu
diff --git a/src/plugins/intel_gpu/src/plugin/transformations/keep_gqa_kv_scale_precision.hpp b/src/plugins/intel_gpu/src/plugin/transformations/keep_gqa_kv_scale_precision.hpp
@@ -0,0 +1,20 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#pragma once
+
+#include "openvino/pass/matcher_pass.hpp"
+
+namespace ov::intel_gpu {
+
+// Marks GroupQueryAttention KV dequant scales (k_scale/v_scale) precision-sensitive so
+// ConvertPrecision keeps them fp32. GPU runs ConvertPrecision before the GQA op is decomposed,
+// and the intact op requires fp32 scales (com.microsoft spec). Applies to a quantized KV cache only.
+class KeepGQAKVScalePrecision : public ov::pass::MatcherPass {
+public:
+    OPENVINO_MATCHER_PASS_RTTI("ov::intel_gpu::KeepGQAKVScalePrecision");
+    KeepGQAKVScalePrecision();
+};
+
+}  // namespace ov::intel_gpu
diff --git a/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp b/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp
@@ -100,6 +100,7 @@
 #include "plugin/transformations/fuse_moe_router_scale.hpp"
 #include "plugin/transformations/increase_position_ids_precision.hpp"
 #include "plugin/transformations/indirect_kv_cache.hpp"
+#include "plugin/transformations/keep_gqa_kv_scale_precision.hpp"
 #include "plugin/transformations/keep_moe_3gemm_const_precision.hpp"
 #include "plugin/transformations/keep_xattention_threshold_precision.hpp"
 #include "plugin/transformations/kv_cache_compression.hpp"
@@ -764,6 +765,9 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
             ov::element::TypeVector{ov::element::i32, ov::element::u32, ov::element::u16}, add_precision_sensitive_convert);
         // Keep xattention threshold in fp32 to avoid boundary issues caused by fp16 quantization.
         manager.register_pass<ov::intel_gpu::KeepXAttentionThresholdPrecision>();
+        // Keep GroupQueryAttention quantized-KV scales fp32 through the ConvertPrecision below
+        // (the intact op requires fp32 scales; it is decomposed later in CommonOptimizations).
+        manager.register_pass<ov::intel_gpu::KeepGQAKVScalePrecision>();
 
         manager.register_pass<ov::pass::ConvertPrecision>(fp_convert_precision_map,
                                                           empty_fuse_map,
diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/scatter_update_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/scatter_update_gpu_test.cpp
@@ -2287,3 +2287,59 @@ TEST(scatter_update_gpu_fp32, d8111_axisB_first_iteration_kernel_check) {
         }
     }
 }
+
+// Integer-typed ScatterUpdate (i8/u8). Exercises the byte-movement path enabled for a quantized KV
+// cache (GroupQueryAttention int8/int4). Dictionary is scattered by index; result is exact.
+template <typename T>
+void test_int_d8111_axisB(bool is_caching_test) {
+    //  Dictionary : 8x1x1x1, Indexes : 4x1x1x1, Updates : 4x1x1x1, Axis : 0
+    //  Indexes:    4, 3, 1, 7
+    //  Updates:    9, 10, 11, 12
+    //  Dictionary: 1, 2, 3, 4, 5, 6, 7, 8
+    //  Output:     1, 11, 3, 10, 9, 6, 7, 12
+    const auto dt = ov::element::from<T>();
+    auto& engine = get_test_engine();
+
+    for (const auto target_format : formats2D) {
+        auto input1 = engine.allocate_memory({dt, plain_2d_format, tensor{8, 1, 1, 1}});               // Dictionary
+        auto input2 = engine.allocate_memory({data_types::f32, plain_2d_format, tensor{4, 1, 1, 1}});  // Indexes
+        auto input3 = engine.allocate_memory({dt, plain_2d_format, tensor{4, 1, 1, 1}});               // Updates
+        auto axis = 0;
+
+        set_values<T>(input1, {T(1), T(2), T(3), T(4), T(5), T(6), T(7), T(8)});
+        set_values(input2, {4.f, 3.f, 1.f, 7.f});
+        set_values<T>(input3, {T(9), T(10), T(11), T(12)});
+
+        topology topology;
+        topology.add(input_layout("InputDictionary", input1->get_layout()));
+        topology.add(input_layout("InputText", input2->get_layout()));
+        topology.add(input_layout("InputUpdates", input3->get_layout()));
+        topology.add(reorder("DictionaryReordered", input_info("InputDictionary"), target_format, dt));
+        topology.add(reorder("TextReordered", input_info("InputText"), target_format, data_types::f32));
+        topology.add(reorder("UpdatesReordered", input_info("InputUpdates"), target_format, dt));
+        topology.add(scatter_update("scatter_update", input_info("DictionaryReordered"), input_info("TextReordered"), input_info("UpdatesReordered"), axis));
+        topology.add(reorder("out", input_info("scatter_update"), plain_2d_format, dt));
+
+        cldnn::network::ptr network = get_network(engine, topology, get_test_default_config(engine), get_test_stream_ptr(), is_caching_test);
+        network->set_input_data("InputDictionary", input1);
+        network->set_input_data("InputText", input2);
+        network->set_input_data("InputUpdates", input3);
+        auto outputs = network->execute();
+
+        auto output = outputs.at("out").get_memory();
+        cldnn::mem_lock<T, mem_lock_type::read> output_ptr(output, get_test_stream());
+
+        std::vector<T> expected_results = {T(1), T(11), T(3), T(10), T(9), T(6), T(7), T(12)};
+        for (size_t i = 0; i < expected_results.size(); ++i) {
+            ASSERT_EQ(expected_results[i], output_ptr[i]) << "i=" << i << ", target_format=" << target_format;
+        }
+    }
+}
+
+TEST(scatter_update_gpu_i8, d8111_axisB) {
+    test_int_d8111_axisB<int8_t>(false);
+}
+
+TEST(scatter_update_gpu_u8, d8111_axisB) {
+    test_int_d8111_axisB<uint8_t>(false);
+}
diff --git a/src/plugins/intel_gpu/tests/unit/test_cases/slice_gpu_test.cpp b/src/plugins/intel_gpu/tests/unit/test_cases/slice_gpu_test.cpp
@@ -331,4 +331,45 @@ TYPED_TEST(SliceTest, bfzyx) {
     this->RunAllTestCasesForParams(params);
 }
 
+// int8 Slice: dedicated small-value case (the shared TYPED_TEST data exceeds the i8 range). Covers the
+// i8 output path enabled for a quantized KV cache (GroupQueryAttention int8/int4 dynamic Slice+Concat).
+TEST(slice_gpu_i8, bfyx) {
+    auto& engine = get_test_engine();
+
+    // input [1,1,2,4] i8, slice axis 2 -> [1,1,1,4] (second row).
+    auto input = engine.allocate_memory({ov::PartialShape{1, 1, 2, 4}, data_types::i8, format::bfyx});
+    auto start = engine.allocate_memory({ov::PartialShape{1}, data_types::i64, format::bfyx});
+    auto stop = engine.allocate_memory({ov::PartialShape{1}, data_types::i64, format::bfyx});
+    auto step = engine.allocate_memory({ov::PartialShape{1}, data_types::i64, format::bfyx});
+    auto axes = engine.allocate_memory({ov::PartialShape{1}, data_types::i64, format::bfyx});
+
+    set_values<int8_t>(input, {1, 2, 3, 4, -5, -6, -7, -8});
+    set_values<int64_t>(start, {1});
+    set_values<int64_t>(stop, {2});
+    set_values<int64_t>(step, {1});
+    set_values<int64_t>(axes, {2});
+
+    topology topology;
+    topology.add(input_layout("input", input->get_layout()));
+    topology.add(data("start", start));
+    topology.add(data("stop", stop));
+    topology.add(data("step", step));
+    topology.add(data("axes", axes));
+    topology.add(slice("slice", {input_info("input"), input_info("start"), input_info("stop"), input_info("step"), input_info("axes")}));
+
+    ExecutionConfig config = get_test_default_config(engine);
+    config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
+    cldnn::network network(engine, topology, config);
+    network.set_input_data("input", input);
+    auto outputs = network.execute();
+
+    auto output = outputs.at("slice").get_memory();
+    cldnn::mem_lock<int8_t, mem_lock_type::read> output_ptr(output, get_test_stream());
+
+    std::vector<int8_t> expected_results = {-5, -6, -7, -8};
+    ASSERT_EQ(output_ptr.size(), expected_results.size());
+    for (size_t i = 0; i < expected_results.size(); ++i)
+        ASSERT_EQ(expected_results[i], output_ptr[i]) << "i=" << i;
+}
+
 } // anonymous namespace
diff --git a/src/plugins/intel_gpu/tests/unit/transformations/keep_gqa_kv_scale_precision_test.cpp b/src/plugins/intel_gpu/tests/unit/transformations/keep_gqa_kv_scale_precision_test.cpp
@@ -0,0 +1,109 @@
+// Copyright (C) 2018-2026 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include "plugin/transformations/keep_gqa_kv_scale_precision.hpp"
+
+#include <gtest/gtest.h>
+
+#include "common_test_utils/ov_test_utils.hpp"
+#include "openvino/core/model.hpp"
+#include "openvino/op/group_query_attention.hpp"
+#include "openvino/op/parameter.hpp"
+#include "openvino/op/result.hpp"
+#include "openvino/op/util/precision_sensitive_attribute.hpp"
+#include "openvino/pass/manager.hpp"
+
+using namespace testing;
+using namespace ov::intel_gpu;
+
+namespace {
+
+using ov::op::internal::GroupQueryAttention;
+
+// k_scale / v_scale are inputs 12 / 13 of GroupQueryAttention.
+constexpr size_t k_scale_idx = 12;
+constexpr size_t v_scale_idx = 13;
+
+// Builds a 14-input GroupQueryAttention model. When quantized, the past KV cache is i8 and the op
+// carries kv_cache_bit_width=8 + a per-channel quant scheme, so is_kv_quantized() is true; otherwise
+// the cache is f16 with the default (unquantized) attributes. k_scale/v_scale (inputs 12/13) are f32
+// in both cases so the no-op test can assert on the exact ports the pass targets.
+std::shared_ptr<ov::Model> make_gqa_model(bool quantized) {
+    const auto kv_et = quantized ? ov::element::i8 : ov::element::f16;
+
+    ov::ParameterVector params(14);
+    params[0] = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::PartialShape{1, 2, 1, 16});  // query
+    params[1] = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::PartialShape{1, 1, 1, 16});  // key
+    params[2] = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::PartialShape{1, 1, 1, 16});  // value
+    params[3] = std::make_shared<ov::op::v0::Parameter>(kv_et, ov::PartialShape{1, 1, 8, 16});             // past_key
+    params[4] = std::make_shared<ov::op::v0::Parameter>(kv_et, ov::PartialShape{1, 1, 8, 16});             // past_value
+    params[5] = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});            // seqlens_k
+    params[6] = std::make_shared<ov::op::v0::Parameter>(ov::element::i32, ov::PartialShape{1});            // total_sequence_length
+    for (size_t i = 7; i <= 11; ++i)  // rotary / optional placeholders, unread by validate_and_infer_types
+        params[i] = std::make_shared<ov::op::v0::Parameter>(ov::element::f16, ov::PartialShape{1});
+    params[k_scale_idx] = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1});  // k_scale
+    params[v_scale_idx] = std::make_shared<ov::op::v0::Parameter>(ov::element::f32, ov::PartialShape{1});  // v_scale
+
+    ov::OutputVector args(params.begin(), params.end());
+    auto gqa = std::make_shared<GroupQueryAttention>(args,
+                                                     /*num_heads*/ 2,
+                                                     /*kv_num_heads*/ 1,
+                                                     /*scale*/ 0.0f,
+                                                     /*do_rotary*/ false,
+                                                     /*rotary_interleaved*/ false,
+                                                     /*kv_cache_bit_width*/ quantized ? 8 : 0,
+                                                     /*k_quant_type*/ quantized ? "PER_CHANNEL" : "NONE",
+                                                     /*v_quant_type*/ quantized ? "PER_CHANNEL" : "NONE");
+
+    ov::ResultVector results{std::make_shared<ov::op::v0::Result>(gqa->output(0)),
+                             std::make_shared<ov::op::v0::Result>(gqa->output(1)),
+                             std::make_shared<ov::op::v0::Result>(gqa->output(2))};
+    return std::make_shared<ov::Model>(results, params, "gqa_model");
+}
+
+std::shared_ptr<GroupQueryAttention> get_gqa(const std::shared_ptr<ov::Model>& model) {
+    return ov::as_type_ptr<GroupQueryAttention>(model->get_results().at(0)->input_value(0).get_node_shared_ptr());
+}
+
+}  // namespace
+
+// Quantized KV cache: the pass must mark k_scale (12) and v_scale (13) as precision-sensitive so that
+// ConvertPrecision keeps them fp32. The model_ref comparison (RUNTIME_KEYS) reliably catches a missing mark.
+TEST_F(TransformationTestsF, KeepGQAKVScalePrecisionMarksQuantizedKVScales) {
+    model = make_gqa_model(/*quantized*/ true);
+
+    model_ref = make_gqa_model(/*quantized*/ true);
+    auto gqa_ref = get_gqa(model_ref);
+    ov::mark_as_precision_sensitive(gqa_ref->input(k_scale_idx));
+    ov::mark_as_precision_sensitive(gqa_ref->input(v_scale_idx));
+
+    manager.register_pass<KeepGQAKVScalePrecision>();
+}
+
+// Explicit positive check, independent of the fixture comparator.
+TEST(KeepGQAKVScalePrecisionTest, QuantizedKVScalesMarked) {
+    auto model = make_gqa_model(/*quantized*/ true);
+    auto gqa = get_gqa(model);
+
+    ov::pass::Manager manager;
+    manager.register_pass<KeepGQAKVScalePrecision>();
+    manager.run_passes(model);
+
+    EXPECT_TRUE(ov::is_precision_sensitive(gqa->input(k_scale_idx)));
+    EXPECT_TRUE(ov::is_precision_sensitive(gqa->input(v_scale_idx)));
+}
+
+// Non-quantized KV cache: is_kv_quantized() is false, so the pass must be a no-op. The directional
+// comparator cannot catch an erroneous mark, so assert the absence explicitly.
+TEST(KeepGQAKVScalePrecisionTest, NonQuantizedKVIsNoOp) {
+    auto model = make_gqa_model(/*quantized*/ false);
+    auto gqa = get_gqa(model);
+
+    ov::pass::Manager manager;
+    manager.register_pass<KeepGQAKVScalePrecision>();
+    manager.run_passes(model);
+
+    EXPECT_FALSE(ov::is_precision_sensitive(gqa->input(k_scale_idx)));
+    EXPECT_FALSE(ov::is_precision_sensitive(gqa->input(v_scale_idx)));
+}
