Refactor: Move transpose check to backend and use standard slicing

Author: Varshith-Yadav

src/nncf/quantization/algorithms/weight_compression/awq.py

@@ -29,8 +29,6 @@
 from nncf.quantization.algorithms.weight_compression.activation_stats import process_stats
 from nncf.quantization.algorithms.weight_compression.backend import WeightCompressionAlgoBackend
 from nncf.quantization.algorithms.weight_compression.config import WeightCompressionParameters
-from nncf.quantization.algorithms.weight_compression.tensor_slicing import get_weight_slice
-from nncf.quantization.algorithms.weight_compression.tensor_slicing import set_weight_slice
 from nncf.quantization.algorithms.weight_compression.weight_lowering import float_quantize_dequantize_weight
 from nncf.quantization.algorithms.weight_compression.weight_lowering import integer_quantize_dequantize_weight
 from nncf.quantization.passes import transform_to_inference_graph
@@ -183,7 +181,9 @@ def apply(
                     prev_weight = self._backend_entity.get_weight(merge_node, prev_weight_port_id, model, graph)
 
                     prev_statistics = statistics[merge_node.node_name]
-                scale = self._data_aware_step(wp, weight, statistics[k], prev_weight, prev_statistics, weight_port_id)
+                scale = self._data_aware_step(
+                    wp, weight, statistics[k], prev_weight, prev_statistics, weight_port_id, graph
+                )
 
             w_scale = fns.unsqueeze(scale, 1 - wp.reduction_axes[0])
             a_scale = fns.unsqueeze(1.0 / scale, wp.reduction_axes[0])
@@ -212,7 +212,9 @@ def apply(
 
         return transformed_model
 
-    def _data_aware_step(self, wp, weight, statistics, prev_weight=None, prev_statistics=None, weight_port_id=None):
+    def _data_aware_step(
+        self, wp, weight, statistics, prev_weight=None, prev_statistics=None, weight_port_id=None, graph=None
+    ):
         alpha_step = (self._alpha_max - self._alpha_min) / self._steps
         config = wp.compression_config
         s, X = process_stats(statistics, self._subset_size)
@@ -222,8 +224,8 @@ def _data_aware_step(self, wp, weight, statistics, prev_weight=None, prev_statis
         assert isinstance(wp.reduction_axes, tuple) and len(wp.reduction_axes) == 1
         reduction_axis = wp.reduction_axes[0]
 
-        # Get transpose_b value to handle weight shape correctly
-        transpose_b = wp.node_with_weight.layer_attributes.constant_attributes[weight_port_id]["transpose"]
+        # Get transpose_b value from backend to handle weight shape correctly in a backend-agnostic way
+        transpose_b = self._backend_entity.get_weight_transpose_b(wp.node_with_weight, weight_port_id, graph)
 
         prev_s, prev_w = None, None
         if prev_statistics is not None and prev_weight is not None:
@@ -244,7 +246,7 @@ def _data_aware_step(self, wp, weight, statistics, prev_weight=None, prev_statis
 
         groups_to_correct = list(groups_to_correct)
 
-        # Remove the old transpose logic - we'll use get_weight_slice instead
+        # Remove the old transpose logic - we now rely on explicit transpose_b handling
         # if reduction_axis == 0:
         #     weight = fns.transpose(weight)
         #     reduction_axis = 1
@@ -263,8 +265,11 @@ def _data_aware_step(self, wp, weight, statistics, prev_weight=None, prev_statis
             a_max = 1e2
             gscale = fns.clip(gscale, a_min=a_min, a_max=a_max)
 
-            # Use get_weight_slice instead of hardcoded slicing
-            gweight = get_weight_slice(weight, slice(offset, offset + group_size), transpose_b)
+            # Slice weight block along the input-channel dimension taking transpose_b into account
+            if transpose_b:
+                gweight = weight[:, offset : offset + group_size]
+            else:
+                gweight = weight[offset : offset + group_size, :]
             gacts = X[offset : offset + group_size, :]
 
             fp32_out = fns.matmul(gweight, gacts)
@@ -310,8 +315,8 @@ def _data_aware_step(self, wp, weight, statistics, prev_weight=None, prev_statis
                 alpha += alpha_step
 
             if best_scale is not None:
-                # Use set_weight_slice for assignment
-                set_weight_slice(scale, slice(offset, offset + group_size), best_scale, transpose_b)
+                # Assign best_scale to the corresponding slice of the scale vector
+                scale[offset : offset + group_size] = best_scale
 
         return scale
 

src/nncf/quantization/algorithms/weight_compression/backend.py

@@ -88,6 +88,20 @@ def get_reduction_axes(node_with_weight: NNCFNode, weight_port_id: int, graph: N
         :return: Reduction shape in tuple format or None if not applicable.
         """
 
+    @staticmethod
+    @abstractmethod
+    def get_weight_transpose_b(node_with_weight: NNCFNode, weight_port_id: int, graph: NNCFGraph) -> bool:
+        """
+        Returns whether the weight input of the given node is treated as transposed (transpose_b=True).
+
+        This is backend-specific and abstracts away how the underlying framework stores MatMul/Gemm attributes.
+
+        :param node_with_weight: The node with weight.
+        :param weight_port_id: The input port ID that corresponds to the weight.
+        :param graph: The model graph.
+        :return: True if the backend treats the weight as transposed (e.g., [Out, In]), False otherwise.
+        """
+
     @staticmethod
     @abstractmethod
     def get_weight_names_and_port_ids(node: NNCFNode, graph: NNCFGraph) -> list[tuple[str, int]]:

src/nncf/quantization/algorithms/weight_compression/gptq.py

@@ -27,8 +27,6 @@
 from nncf.quantization.algorithms.weight_compression.config import WeightCompressionParameters
 from nncf.quantization.algorithms.weight_compression.parameters import CompressedWeight
 from nncf.quantization.algorithms.weight_compression.scale_estimation import ScaleEstimation
-from nncf.quantization.algorithms.weight_compression.tensor_slicing import get_weight_slice
-from nncf.quantization.algorithms.weight_compression.tensor_slicing import set_weight_slice
 from nncf.quantization.algorithms.weight_compression.weight_lowering import calculate_float_quantization_params
 from nncf.quantization.algorithms.weight_compression.weight_lowering import calculate_integer_quantization_params
 from nncf.quantization.algorithms.weight_compression.weight_lowering import float_quantize_dequantize_weight
@@ -220,16 +218,19 @@ def _quantize_weights(
         )
         weight_tensor = fns.astype(weight_tensor, TensorDataType.float32)
 
-        # Get transpose_b value to handle weight shape correctly
-        transpose_b = wc_params.node_with_weight.layer_attributes.constant_attributes[wc_params.weight_port_id][
-            "transpose"
-        ]
+        # Get transpose_b value via backend to handle weight shape correctly in a backend-agnostic way
+        transpose_b = self._backend_entity.get_weight_transpose_b(
+            wc_params.node_with_weight, wc_params.weight_port_id, graph
+        )
 
         dead_indices = fns.diag(hessian) == 0
         hessian[dead_indices, dead_indices] = 1
 
-        # Zero out dead indices using utility helper
-        set_weight_slice(weight_tensor, dead_indices, 0, transpose_b)
+        # Zero out dead indices along the input-channel dimension
+        if transpose_b:
+            weight_tensor[:, dead_indices] = 0
+        else:
+            weight_tensor[dead_indices, :] = 0
 
         scales = []
         zero_points = []
@@ -260,26 +261,38 @@ def _quantize_weights(
             i2 = min(i1 + self._block_size, columns)
             count = i2 - i1
 
-            # Extract weight block using utility helper
-            weight_block = get_weight_slice(weight_tensor, slice(i1, i2), transpose_b).clone()
+            # Extract weight block along the input-channel dimension
+            if transpose_b:
+                weight_block = weight_tensor[:, i1:i2].clone()
+            else:
+                weight_block = weight_tensor[i1:i2, :].clone()
             quantized_block = fns.zeros_like(weight_block)
             error_block = fns.zeros_like(weight_block)
             loss_block = fns.zeros_like(weight_block)
             hessian_inv_block = hessian_inv[i1:i2, i1:i2]
 
             for i in range(count):
-                weight_col = get_weight_slice(weight_block, i, transpose_b)
+                if transpose_b:
+                    weight_col = weight_block[:, i]
+                else:
+                    weight_col = weight_block[i, :]
                 hessian_diag_val = hessian_inv_block[i, i]
 
                 if (i1 + i) % group_size == 0:
                     if not block_compression_config.is_integer:
-                        weight_slice = get_weight_slice(weight_tensor, slice(i1 + i, i1 + i + group_size), transpose_b)
+                        if transpose_b:
+                            weight_slice = weight_tensor[:, i1 + i : i1 + i + group_size]
+                        else:
+                            weight_slice = weight_tensor[i1 + i : i1 + i + group_size, :]
                         scale = calculate_float_quantization_params(
                             weight_slice, reduction_axes, block_compression_config
                         )
                         scales.append(scale)
                     else:
-                        weight_slice = get_weight_slice(weight_tensor, slice(i1 + i, i1 + i + group_size), transpose_b)
+                        if transpose_b:
+                            weight_slice = weight_tensor[:, i1 + i : i1 + i + group_size]
+                        else:
+                            weight_slice = weight_tensor[i1 + i : i1 + i + group_size, :]
                         if self._scale_estimation and block_compression_config.num_bits == 4:
                             activations = [inp[..., (i1 + i) : (i1 + i + group_size)] for inp in inputs]
                             wc_statistics = ScaleEstimation.activations_to_wc_statistics(activations)
@@ -312,24 +325,34 @@ def _quantize_weights(
                         precomputed_zero_point=zero_points[-1],
                     )
                 quantized_col = fns.flatten(quantized_col)
-                set_weight_slice(quantized_block, i, quantized_col, transpose_b)
+                if transpose_b:
+                    quantized_block[:, i] = quantized_col
+                else:
+                    quantized_block[i, :] = quantized_col
                 loss_col = (weight_col - quantized_col) ** 2 / hessian_diag_val**2
-                set_weight_slice(loss_block, i, loss_col, transpose_b)
+                if transpose_b:
+                    loss_block[:, i] = loss_col
+                else:
+                    loss_block[i, :] = loss_col
 
                 error_col = (weight_col - quantized_col) / hessian_diag_val
                 if transpose_b:
                     weight_block[:, i:] -= fns.matmul(
                         fns.unsqueeze(error_col, 1), fns.unsqueeze(hessian_inv_block[i, i:], 0)
                     )
-                    set_weight_slice(error_block, i, error_col, transpose_b)
+                    error_block[:, i] = error_col
                 else:
                     weight_block[i:, :] -= fns.matmul(
                         fns.unsqueeze(error_col, 0), fns.unsqueeze(hessian_inv_block[i:, i], 1)
                     )
-                    set_weight_slice(error_block, i, error_col, transpose_b)
+                    error_block[i, :] = error_col
 
-            set_weight_slice(quantized_tensor, slice(i1, i2), quantized_block, transpose_b)
-            set_weight_slice(losses, slice(i1, i2), loss_block / 2, transpose_b)
+            if transpose_b:
+                quantized_tensor[:, i1:i2] = quantized_block
+                losses[:, i1:i2] = loss_block / 2
+            else:
+                quantized_tensor[i1:i2, :] = quantized_block
+                losses[i1:i2, :] = loss_block / 2
 
             # Update remaining weights with error propagation
             if transpose_b:

src/nncf/quantization/algorithms/weight_compression/lora_correction.py

@@ -14,6 +14,7 @@
 import pandas as pd
 
 import nncf
+from nncf.common.graph import NNCFGraph
 from nncf.common.logging import nncf_logger
 from nncf.common.utils.debug import DEBUG_LOG_DIR
 from nncf.common.utils.debug import is_debug
@@ -108,24 +109,29 @@ def is_applicable(self, wc_params: WeightCompressionParameters):
         return wc_params.compression_config.num_bits == 4
 
     def calculate_adapters(
-        self, weight: Tensor, compressed_weight: CompressedWeight, wc_params: WeightCompressionParameters
+        self,
+        weight: Tensor,
+        compressed_weight: CompressedWeight,
+        wc_params: WeightCompressionParameters,
+        graph: NNCFGraph,
     ) -> tuple[Tensor, Tensor, list[float]]:
         """
         Calculates low rank matrices for a given original and compressed weights.
 
         :param weight: original floating-point weight matrix.
         :param compressed_weight: compressed weight matrix.
         :param wc_params: parameters of weight compression.
+        :param graph: The model graph.
         :return: two low rank matrices in the order of execution of corresponding linear layers.
         """
         layer_name = wc_params.node_with_weight.node_name
         layer_statistics = self._statistics[layer_name]
         is_debug = self._debug_interface is not None
 
-        # Get transpose_b value to handle weight shape correctly
-        transpose_b = wc_params.node_with_weight.layer_attributes.constant_attributes[wc_params.weight_port_id][
-            "transpose"
-        ]
+        # Get transpose_b value via backend to handle weight shape correctly in a backend-agnostic way
+        transpose_b = self._backend_entity.get_weight_transpose_b(
+            wc_params.node_with_weight, wc_params.weight_port_id, graph
+        )
 
         lora_A, lora_B, mean_noises = self.calculate_low_rank_matrices(
             weight,

src/nncf/quantization/algorithms/weight_compression/onnx_backend.py

@@ -149,6 +149,26 @@ def get_reduction_axes(node_with_weight: NNCFNode, weight_port_id: int, graph: N
             channel_axes = (0,) + channel_axes
         return get_reduction_axes(channel_axes, const_shape)
 
+    @staticmethod
+    def get_weight_transpose_b(node_with_weight: NNCFNode, weight_port_id: int, graph: NNCFGraph) -> bool:
+        """
+        Returns the equivalent of transpose_b for ONNX MatMul/Gemm nodes.
+
+        For MatMul the initializer layout already follows the expected [K, N] contract of the op,
+        so we treat weights as transposed (transpose_b=True) by default. For Gemm we respect the
+        transB attribute when the corresponding input port matches the B input.
+        """
+        # Gemm-specific handling: attribute name and semantics follow ONNX Gemm spec.
+        if node_with_weight.metatype is onnx_metatypes.ONNXGemmMetatype and weight_port_id == 1:
+            node_attrs = node_with_weight.layer_attributes.node_attrs
+            # In Gemm, transB=1 means the B input is transposed.
+            trans_b_attr = node_attrs.get("transB", 0)
+            return bool(trans_b_attr)
+
+        # For MatMul and other ops, rely on the fact that the stored initializer already matches
+        # the expected layout for the backend kernels and treat it as transpose_b=True.
+        return True
+
     @staticmethod
     def target_point(target_type: TargetType, target_node_name: str, port_id: int) -> ONNXTargetPoint:
         return ONNXTargetPoint(target_type, target_node_name, port_id)

src/nncf/quantization/algorithms/weight_compression/openvino_backend.py

@@ -103,6 +103,17 @@ def get_reduction_axes(node_with_weight: NNCFNode, weight_port_id: int, graph: N
         const_shape = node_with_weight.layer_attributes.constant_attributes[weight_port_id]["shape"]
         return get_reduction_axes(channel_axes, const_shape)
 
+    @staticmethod
+    def get_weight_transpose_b(node_with_weight: NNCFNode, weight_port_id: int, graph: NNCFGraph) -> bool:
+        """
+        Returns the value of the transpose_b attribute for the given weight input.
+
+        OpenVINO stores this information in the constant_attributes of the node layer attributes.
+        Default to True for backward compatibility with older models that do not expose the flag.
+        """
+        const_attrs = node_with_weight.layer_attributes.constant_attributes[weight_port_id]
+        return const_attrs.get("transpose", True)
+
     @staticmethod
     def target_point(target_type: TargetType, target_node_name: str, port_id: int) -> OVTargetPoint:
         return OVTargetPoint(target_type, target_node_name, port_id)
@@ -381,7 +392,7 @@ def transform_model(
                 compressed_weight.tensor = compressed_weight.tensor.as_numpy_tensor()
                 if compressed_weight.zero_point is not None:
                     compressed_weight.zero_point = compressed_weight.zero_point.as_numpy_tensor()
-                adapters = lora_correction_algo.calculate_adapters(weight, compressed_weight, wc_params)
+                adapters = lora_correction_algo.calculate_adapters(weight, compressed_weight, wc_params, graph)
                 self.insert_adapters(wc_params, *adapters, int8_lora=lora_correction_algo.use_int8_adapters)
         self.name_to_node_mapping = None
 

src/nncf/quantization/algorithms/weight_compression/scale_estimation.py

@@ -141,8 +141,8 @@ def apply(
 
             weight = self._backend_entity.get_weight(wp.node_with_weight, weight_port_id, model, graph)
 
-            # Get transpose_b value to handle weight shape correctly
-            transpose_b = wp.node_with_weight.layer_attributes.constant_attributes[weight_port_id]["transpose"]
+            # Get transpose_b value via backend to handle weight shape correctly in a backend-agnostic way
+            transpose_b = self._backend_entity.get_weight_transpose_b(wp.node_with_weight, weight_port_id, graph)
 
             scale, zero_point = self.calculate_quantization_params(
                 stats,