[GPU] Add UINT2 (2-bit) weight compression support via oneDNN fully_connected

src/plugins/intel_gpu/src/graph/debug_helper.cpp

@@ -318,6 +318,31 @@ void log_memory_to_file(memory::ptr mem, layout data_layout, stream& stream, std
         dump<uint8_t>(actual_mem, stream, file_stream, dump_raw);
     } else if (mem_dt == cldnn::data_types::i4 || mem_dt == cldnn::data_types::u4) {
         dump_i4u4(mem_dt, actual_mem, stream, file_stream, dump_raw);
+    } else if (mem_dt == cldnn::data_types::u2) {
+        // u2: 4 elements per byte, dump each 2-bit value
+        auto&& size = actual_mem->get_layout().get_tensor();
+        file_stream << "shape: " << size.to_string() << " ";
+        file_stream << "(count: " << size.count()
+                    << ", original format: " << cldnn::fmt_to_str(actual_mem->get_layout().format) << ")"
+                    << (dump_raw ? " raw data" : "") << std::endl;
+        if (size.count() == 0) {
+            file_stream << "Empty buffer" << std::endl;
+        } else if (!dump_raw) {
+            // Non-raw mode would need layout/padding-aware indexing; not supported yet, same as dump_i4u4.
+            GPU_DEBUG_COUT << " supports raw dump only" << std::endl;
+        } else {
+            mem_lock<uint8_t, mem_lock_type::read> lock(actual_mem, stream);
+            auto* mem_ptr = lock.data();
+            std::stringstream buffer;
+            for (size_t i = 0; i < lock.size(); ++i) {
+                uint8_t byte = mem_ptr[i];
+                buffer << static_cast<int>((byte >> 0) & 0x3) << std::endl;
+                buffer << static_cast<int>((byte >> 2) & 0x3) << std::endl;
+                buffer << static_cast<int>((byte >> 4) & 0x3) << std::endl;
+                buffer << static_cast<int>((byte >> 6) & 0x3) << std::endl;
+            }
+            file_stream << buffer.str();
+        }
     } else
         GPU_DEBUG_COUT << "Dump for this data type is not supported: " << dt_to_str(mem_dt) << std::endl;
 }

src/plugins/intel_gpu/src/graph/impls/onednn/fully_connected_onednn.cpp

@@ -23,22 +23,6 @@ struct fully_connected_onednn : typed_primitive_onednn_impl<fully_connected> {
     static constexpr int COMMON = 0;
     static constexpr int PER_OC = 2;
 
-    // Group (IFM) axis is last for [N, groups] (transpose_b) and second-to-last for [groups, N].
-    static int64_t get_ifm_dim_idx(const layout& param_layout, bool transpose_b_param, size_t weights_rank) {
-        const auto& pshape = param_layout.get_partial_shape();
-        const auto rank = static_cast<int64_t>(pshape.size());
-        // 1D [N] holds only OC; point past it so get_dim() hits the padded 1 (one IFM group).
-        if (rank <= 1)
-            return rank;
-        const auto logical_rank = std::min<int64_t>(rank, std::max<int64_t>(2, static_cast<int64_t>(weights_rank)));
-        // Extra dims are either legacy trailing padding (all 1) or leading broadcast dims.
-        const bool trailing_padded = std::all_of(pshape.begin() + logical_rank, pshape.end(), [](const ov::Dimension& d) {
-            return d.is_static() && d.get_length() == 1;
-        });
-        const int64_t base = trailing_padded ? 0 : rank - logical_rank;
-        return std::max<int64_t>(0, base + (transpose_b_param ? logical_rank - 1 : logical_rank - 2));
-    }
-
     DECLARE_OBJECT_TYPE_SERIALIZATION(cldnn::onednn::fully_connected_onednn)
 
 private:
@@ -69,8 +53,7 @@ struct fully_connected_onednn : typed_primitive_onednn_impl<fully_connected> {
         if (prim->compressed_weights) {
             const auto weights_dt = instance.get_input_layout(1).data_type;
             auto weight_bitwidth = ov::element::Type(weights_dt).bitwidth();
-            OPENVINO_ASSERT(weight_bitwidth == 8 || weight_bitwidth == 4 || weight_bitwidth == 3,
-                            "[GPU] oneDNN supports only 3bit/4bit/8bit compressed weights");
+            OPENVINO_ASSERT(one_of(weight_bitwidth, {8, 4, 3, 2}), "[GPU] oneDNN supports only 2bit/3bit/4bit/8bit compressed weights");
             int idx = prim->bias.is_valid() ? 3 : 2;
 
             if (prim->decompression_scale.is_valid()) {
@@ -181,7 +164,7 @@ struct fully_connected_onednn : typed_primitive_onednn_impl<fully_connected> {
         dnnl::memory::desc weights_md;
         if (weights_layout.data_padding
             && format::is_default_format(weights_layout.format)
-            && (weights_layout.data_type == data_types::i4 || weights_layout.data_type == data_types::u4)) {
+            && one_of(weights_layout.data_type, {data_types::i4, data_types::u4, data_types::u2})) {
             weights_md = onednn::layout_to_memory_desc_strides(weights_layout, weights_fmt);
         } else {
             weights_md = onednn::layout_to_memory_desc(weights_layout, weights_fmt);
@@ -285,11 +268,7 @@ struct fully_connected_onednn : typed_primitive_onednn_impl<fully_connected> {
 
         const kernel_impl_params* impl_params = reinterpret_cast<kernel_impl_params*>(ib.getKernelImplParams());
         auto prim = impl_params->typed_desc<fully_connected>();
-        auto weights_layout = impl_params->get_input_layout(1);
-        const auto& weight_shape = weights_layout.get_partial_shape();
-        auto weight_rank = std::count_if(weight_shape.begin(), weight_shape.end(),
-            [](ov::Dimension d) { return d.get_length() > 1; });
-        weight_rank = std::max(static_cast<int64_t>(2), weight_rank);
+        const auto weight_rank = get_compressed_weights_rank(impl_params->get_input_layout(1));
 
         auto shift_size = std::max<size_t>(prim->input_size - 2, 0);
         const auto& arg = impl_params->get_program().get_node(impl_params->desc->id).as<fully_connected>();
@@ -393,9 +372,7 @@ struct fully_connected_onednn : typed_primitive_onednn_impl<fully_connected> {
             }
 
             auto weights_layout = impl_params.get_input_layout(1);
-            auto weight_shape = weights_layout.get_partial_shape();
-            auto weight_rank = std::count_if(weight_shape.begin(), weight_shape.end(), [](ov::Dimension d) { return d.get_length() > 1; });
-            weight_rank = std::max(static_cast<int64_t>(2), weight_rank);
+            const auto weight_rank = get_compressed_weights_rank(weights_layout);
             OPENVINO_ASSERT(weight_rank <= 3, "Currently only weights with equal to or less than 3D is supported");
             auto shift_size = std::max<size_t>(prim->input_size - 2, 0);
             int per_oc = PER_OC << shift_size;
@@ -405,13 +382,10 @@ struct fully_connected_onednn : typed_primitive_onednn_impl<fully_connected> {
                 auto decompression_scale_idx = ++idx;
                 auto scale_layout = arg.get_dependency(decompression_scale_idx).get_output_layout();
                 ds_data_type = convert_data_type(scale_layout.data_type);
-                // IFM (K) dimension position depends on weight layout orientation.
-                const auto ifm_dim_idx = prim->weights_transposed ? (weight_rank - 1) : (weight_rank - 2);
-                const auto ifm = arg.get_dependency(1).get_output_layout().get_dim(ifm_dim_idx);
-                const auto scale_ifm_dim_idx = get_ifm_dim_idx(scale_layout, prim->transpose_b_scale, weight_rank);
-                const auto ngroups = scale_layout.get_dim(scale_ifm_dim_idx);
-                group_size = static_cast<int>(ifm / ngroups);
-                OPENVINO_ASSERT((group_size == 1 || ngroups == 1 || group_size % 16 == 0),
+                const auto groups = get_scale_groups(weights_layout, scale_layout, prim->weights_transposed, prim->transpose_b_scale);
+                const auto ngroups = groups.first;
+                group_size = static_cast<int>(groups.second);
+                OPENVINO_ASSERT(is_supported_scale_group(ngroups, group_size),
                     "[GPU] group_size should be aligned to 16 if it is not a single scale group or the group_size is not one.");
                 if (scale_layout.count() == 1) {
                     attr->set_scales(DNNL_ARG_WEIGHTS, COMMON, dnnl::memory::dims{}, ds_data_type);

src/plugins/intel_gpu/src/graph/impls/onednn/fully_connected_onednn.hpp

@@ -69,8 +69,9 @@ struct FullyConnectedImplementationManager : public ImplementationManager {
                          one_of(out_dt, {data_types::f16, data_types::bf16, data_types::f32, data_types::i32, data_types::i8, data_types::u8});
         bool compressed_case = fc_prim->compressed_weights &&
                                one_of(in0_dt, {data_types::f16, data_types::bf16, data_types::f32, data_types::i8, data_types::u8}) &&
-                               one_of(wei_dt, {data_types::u8, data_types::i8, data_types::u4, data_types::i4, data_types::u3}) &&
-                               one_of(out_dt, {data_types::f16, data_types::bf16, data_types::f32, data_types::u8, data_types::i8});
+                               one_of(wei_dt, {data_types::u8, data_types::i8, data_types::u4, data_types::i4, data_types::u3, data_types::u2}) &&
+                               one_of(out_dt, {data_types::f16, data_types::bf16, data_types::f32, data_types::u8, data_types::i8}) &&
+                               (wei_dt != data_types::u2 || in0_dt == data_types::f16);  // u2 oneDNN kernel requires f16 activations
         const bool fp_compressed_case = fc_prim->compressed_weights && one_of(in0_dt, {data_types::f8e4m3, data_types::f4e2m1, data_types::f8e5m2}) &&
                                         one_of(wei_dt, {data_types::f8e4m3, data_types::f4e2m1, data_types::f8e5m2}) && one_of(out_dt, {data_types::f16, data_types::f32});
         if (!f16f16_case && !bf16bf16_case && !f32f32_case && !u8s8_case && !compressed_case && !fp_compressed_case) {
@@ -81,13 +82,25 @@ struct FullyConnectedImplementationManager : public ImplementationManager {
             if (fc_prim->decompression_zero_point.is_valid()) {
                 const auto decompression_zp_idx = fc_prim->bias.is_valid() ? 4 : 3;
                 const auto decompression_zp_dt = fc_node.get_input_layout(decompression_zp_idx).data_type;
-                const bool is_wei_compressed = one_of(wei_dt, {ov::element::Type_t::i4, ov::element::Type_t::u4, ov::element::Type_t::u3, ov::element::Type_t::u8});
+                const bool is_wei_compressed = one_of(wei_dt, {ov::element::Type_t::i4, ov::element::Type_t::u4, ov::element::Type_t::u3, ov::element::Type_t::u8,
+                                                               ov::element::Type_t::u2});
+                // u2 zero-point excluded: oneDNN's u2 kernel cannot unpack a packed-u2 zero-point tensor
                 const bool is_zp_compressed = one_of(decompression_zp_dt, {ov::element::Type_t::i4, ov::element::Type_t::u4,
                                                     ov::element::Type_t::u8, ov::element::Type_t::i8});
                 if (!is_wei_compressed || !is_zp_compressed) {
                     LOG_AND_RETURN_FALSE(node);
                 }
             }
+
+            // Reject u2 scale grouping oneDNN can't handle here, so another FC impl is selected instead.
+            if (wei_dt == data_types::u2 && fc_prim->decompression_scale.is_valid()) {
+                const auto& scale_layout = fc_node.get_input_layout(fc_prim->bias.is_valid() ? 3 : 2);
+                const auto groups = get_scale_groups(fc_node.weights().get_output_layout(false), scale_layout,
+                                                     fc_prim->weights_transposed, fc_prim->transpose_b_scale);
+                if (!is_supported_scale_group(groups.first, groups.second)) {
+                    LOG_AND_RETURN_FALSE(node);
+                }
+            }
         }
 
         if (fp_compressed_case && fc_prim->decompression_scale.is_valid() && fc_prim->activation_scale.is_valid()) {

src/plugins/intel_gpu/src/graph/impls/onednn/utils.cpp

@@ -6,6 +6,7 @@
 
 #include <oneapi/dnnl/dnnl_debug.h>
 
+#include <algorithm>
 #include <numeric>
 
 namespace cldnn {
@@ -119,6 +120,43 @@ dnnl::memory::format_tag get_default_data_format(const cldnn::layout& l) {
     }
 }
 
+int64_t get_ifm_dim_idx(const layout& param_layout, bool transpose_b_param, size_t weights_rank) {
+    const auto& pshape = param_layout.get_partial_shape();
+    const auto rank = static_cast<int64_t>(pshape.size());
+    // 1D [N] holds only OC; point past it so get_dim() hits the padded 1 (one IFM group).
+    if (rank <= 1)
+        return rank;
+    const auto logical_rank = std::min<int64_t>(rank, std::max<int64_t>(2, static_cast<int64_t>(weights_rank)));
+    // Extra dims are either legacy trailing padding (all 1) or leading broadcast dims.
+    const bool trailing_padded = std::all_of(pshape.begin() + logical_rank, pshape.end(), [](const ov::Dimension& d) {
+        return d.is_static() && d.get_length() == 1;
+    });
+    const int64_t base = trailing_padded ? 0 : rank - logical_rank;
+    return std::max<int64_t>(0, base + (transpose_b_param ? logical_rank - 1 : logical_rank - 2));
+}
+
+int64_t get_compressed_weights_rank(const layout& weights_layout) {
+    const auto& shape = weights_layout.get_partial_shape();
+    const auto rank = std::count_if(shape.begin(), shape.end(), [](const ov::Dimension& d) {
+        return d.get_length() > 1;
+    });
+    return std::max<int64_t>(2, rank);
+}
+
+std::pair<int64_t, int64_t> get_scale_groups(const layout& weights_layout,
+                                             const layout& scale_layout,
+                                             bool weights_transposed,
+                                             bool transpose_b_scale) {
+    const auto weights_rank = get_compressed_weights_rank(weights_layout);
+    const int64_t ifm = weights_layout.get_dim(weights_transposed ? weights_rank - 1 : weights_rank - 2);
+    const int64_t ngroups = scale_layout.get_dim(get_ifm_dim_idx(scale_layout, transpose_b_scale, weights_rank));
+    return {ngroups, ngroups != 0 ? ifm / ngroups : ifm};
+}
+
+bool is_supported_scale_group(int64_t ngroups, int64_t group_size) {
+    return group_size == 1 || ngroups == 1 || group_size % 16 == 0;
+}
+
 dnnl::memory::format_tag convert_gemm_data_format(dnnl::memory::dims dims, format target) {
     if (dims.size() == target.dimension()) {
         auto tag = convert_data_format(target);
@@ -172,6 +210,7 @@ dnnl::memory::data_type convert_data_type(cldnn::data_types dt) {
         case cldnn::data_types::i4: return dnnl::memory::data_type::s4;
         case cldnn::data_types::u4: return dnnl::memory::data_type::u4;
         case cldnn::data_types::u3: return dnnl::memory::data_type::u3;
+        case cldnn::data_types::u2: return dnnl::memory::data_type::u2;
         case cldnn::data_types::f4e2m1: return dnnl::memory::data_type::f4_e2m1;
         case cldnn::data_types::f8e4m3: return dnnl::memory::data_type::f8_e4m3;
         case cldnn::data_types::f8e5m2: return dnnl::memory::data_type::f8_e5m2;
@@ -288,6 +327,8 @@ int64_t get_offset(const cldnn::layout& l, dnnl::memory::desc&& desc) {
     }
 
     switch (desc.get_data_type()) {
+        case dnnl::memory::data_type::u2:
+            return offset / 4;
         case dnnl::memory::data_type::u3:
             OPENVINO_ASSERT(offset % 8 == 0, "[GPU] u3 memory offset must be byte aligned, but got element offset ", offset);
             return offset * 3 / 8;

src/plugins/intel_gpu/src/graph/impls/onednn/utils.hpp

@@ -14,6 +14,7 @@
 #include "intel_gpu/primitives/reorder.hpp"
 #include "intel_gpu/runtime/format.hpp"
 #include <data_inst.h>
+#include <utility>
 
 namespace cldnn {
 namespace onednn {
@@ -33,6 +34,23 @@ cldnn::format convert_data_format(dnnl::memory::format_tag fmt);
 dnnl::memory::format_tag get_default_data_format(const cldnn::layout& l);
 dnnl::memory::format_tag convert_gemm_data_format(dnnl::memory::dims dims, format target);
 
+// Group (IFM) axis is last for [N, groups] (transpose_b) and second-to-last for [groups, N].
+// Shared between implementation-selection (validate_impl) and primitive creation, so both agree
+// on which axis holds the per-group scale/zero-point count.
+int64_t get_ifm_dim_idx(const layout& param_layout, bool transpose_b_param, size_t weights_rank);
+
+// Rank of compressed weights, ignoring size-1 dims; at least 2.
+int64_t get_compressed_weights_rank(const layout& weights_layout);
+
+// Returns {ngroups, group_size} of the weight-decompression scale along IFM.
+std::pair<int64_t, int64_t> get_scale_groups(const layout& weights_layout,
+                                             const layout& scale_layout,
+                                             bool weights_transposed,
+                                             bool transpose_b_scale);
+
+// oneDNN grouped weight scales need group_size == 1, a single group, or group_size % 16 == 0.
+bool is_supported_scale_group(int64_t ngroups, int64_t group_size);
+
 dnnl::memory::desc layout_to_memory_desc(
     const cldnn::layout& l,
     dnnl::memory::format_tag target_fmt = dnnl::memory::format_tag::undef);

src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp

@@ -1897,6 +1897,12 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
                         return true;
                     }
                 }
+                // u2 oneDNN kernel expects fp16 activations; disabled until i8 activation support is added
+                if (root->get_input_element_type(1) == ov::element::u2) {
+                    GPU_DEBUG_TRACE << root->get_friendly_name() << "  dyn_quan is turned off: u2 weights require fp16 activations" << std::endl;
+                    return true;
+                }
+
                 uint64_t adj_group_size = dynamic_quantization_group_size;
                 const bool is_wei_i8u8 = cldnn::one_of(root->get_input_element_type(1), {ov::element::i8, ov::element::u8});
                 if (ov::intel_gpu::DynamicQuantizeFullyConnected::ShouldUseGs128(is_wei_i8u8, use_gs128_for_int8_per_token, adj_group_size, use_gs128_for_linear_attention)) {