Fix ONNX AWQ-Lite get_weight_scale mixing rows from different blocks

CHANGELOG.rst

@@ -27,6 +27,7 @@ Changelog
 - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11.
 - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit.
 - Fix ONNX quantization with the default ``concat_elimination`` pass crashing in calibration with ``KeyError: '<tensor>_ReduceMin'`` when a ``Concat`` input is a constant, such as the class token of a ViT.
+- Fix ONNX AWQ-Lite's ``get_weight_scale`` computing per-block weight scales from the wrong rows: reshaping the weight directly into ``(block_size, -1)`` interleaves rows from different blocks instead of grouping each block's ``block_size`` consecutive rows together, silently skewing the search scale AWQ-Lite picks for INT4 PTQ.
 
 0.48.0 (2026-10-xx)
 ^^^^^^^^^^^^^^^^^^^

modelopt/onnx/quantization/int4.py

@@ -747,10 +747,16 @@ def get_weight_scale(weight, block_size=None):
             slice_after_padding = slice(org_shape[0])
             weight = _pad(weight, block_size)
             org_shape = weight.shape
-        weight = weight.reshape(block_size, -1)
-    weight_abs_amax = weight.__abs__().max(axis=0, keepdims=True)
-    scale = weight.__abs__() / (weight_abs_amax + np.finfo(weight_abs_amax.dtype).tiny)
-    scale = scale.reshape(org_shape)
+        # Blocks are `block_size` consecutive rows; splitting dim 0 into
+        # (n_blocks, block_size) keeps each block's rows together, unlike
+        # reshape(block_size, -1), which interleaves rows from different blocks.
+        weight = weight.reshape(org_shape[0] // block_size, block_size, org_shape[1])
+        weight_abs_amax = weight.__abs__().max(axis=1, keepdims=True)
+        scale = weight.__abs__() / (weight_abs_amax + np.finfo(weight_abs_amax.dtype).tiny)
+        scale = scale.reshape(org_shape)
+    else:
+        weight_abs_amax = weight.__abs__().max(axis=0, keepdims=True)
+        scale = weight.__abs__() / (weight_abs_amax + np.finfo(weight_abs_amax.dtype).tiny)
     if slice_after_padding is not None:
         scale = scale[slice_after_padding, ...]
     scale = scale.mean(1)

tests/unit/onnx/quantization/test_int4.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+from modelopt.onnx.quantization.int4 import get_weight_scale
+
+
+def _reference_weight_scale(weight: np.ndarray, block_size: int) -> np.ndarray:
+    """Per-block-max scale using consecutive, non-overlapping row blocks."""
+    ci, co = weight.shape
+    n_blocks = -(-ci // block_size)  # ceil division, matches the padding in get_weight_scale
+    padded = np.pad(weight, ((0, n_blocks * block_size - ci), (0, 0)))
+    scale = np.zeros_like(padded)
+    for b in range(n_blocks):
+        block = padded[b * block_size : (b + 1) * block_size]
+        amax = np.abs(block).max(axis=0, keepdims=True)
+        scale[b * block_size : (b + 1) * block_size] = np.abs(block) / (
+            amax + np.finfo(amax.dtype).tiny
+        )
+    return scale[:ci].mean(1)
+
+
+def test_get_weight_scale_groups_consecutive_rows_into_blocks():
+    """A block's max must come from its own rows, not rows from a different block.
+
+    Reshaping (ci, co) directly into (block_size, -1) does not slice `ci` into
+    contiguous blocks of `block_size` rows -- it interleaves rows from every block,
+    so a row's scale ends up computed against the wrong max.
+    """
+    weight = np.zeros((8, 1))
+    weight[0, 0] = 10.0
+    weight[1, 0] = 5.0
+
+    scale = get_weight_scale(weight, block_size=4)
+
+    # Rows 0 and 1 are both in the first 4-row block, whose max is 10.
+    assert scale[0] == 1.0
+    assert scale[1] == 0.5
+    np.testing.assert_array_equal(scale[2:], 0.0)
+
+
+def test_get_weight_scale_matches_reference_when_divisible():
+    rng = np.random.RandomState(0)
+    weight = rng.randn(16, 3)
+    block_size = 4
+
+    scale = get_weight_scale(weight.copy(), block_size)
+
+    np.testing.assert_allclose(scale, _reference_weight_scale(weight, block_size))
+
+
+def test_get_weight_scale_matches_reference_with_padding():
+    rng = np.random.RandomState(1)
+    weight = rng.randn(10, 3)  # not a multiple of block_size
+    block_size = 4
+
+    scale = get_weight_scale(weight.copy(), block_size)
+
+    assert scale.shape == (10,)
+    np.testing.assert_allclose(scale, _reference_weight_scale(weight, block_size))
+
+
+def test_get_weight_scale_without_block_size_is_unchanged():
+    rng = np.random.RandomState(2)
+    weight = rng.randn(16, 3)
+
+    scale = get_weight_scale(weight.copy(), block_size=None)
+
+    amax = np.abs(weight).max(axis=0, keepdims=True)
+    expected = (np.abs(weight) / (amax + np.finfo(amax.dtype).tiny)).mean(1)
+    np.testing.assert_allclose(scale, expected)