GatherBlockQuantized CUDA kernel uses the flat index for the block id, giving wrong results when the quantized axis is not a multiple of block_size

Author: tianleiwuCreated Sep 16, 2026Updated Sep 17, 2026

Labelsquantizationep:CUDA

Describe the issue

The CUDA implementation of com.microsoft.GatherBlockQuantized computes the block (scale / zero-point) index by dividing the flat data index by block_size. That is only correct when the quantized axis is an exact multiple of block_size. When it is not, elements are dequantized with the scale of a neighbouring block, and because blockwise scales are signed the resulting values are frequently sign-flipped, not merely imprecise.

The CPU kernel handles this correctly, so CPU and CUDA disagree on the same model. The op schema explicitly permits a non-multiple: shape inference accepts a scales dimension of ceil(data_dim * components / block_size).

This is a silent wrong-results bug — no error is raised.

https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu#L66

cpp

int64_t in_idx = idx_before * gather_axis_dim * after_gather_dim + idx_at_g * after_gather_dim + idx_after;

int64_t block_id = in_idx / block_size;   // <-- flat index / block_size

in_idx is a flat index into the data tensor, but blocks do not tile the flat array contiguously: each row along the quantized axis starts a fresh group of ceil(K / block_size) blocks. So the correct index is

block_id = row * ceil(K / block_size) + (col / block_size)

which equals in_idx / block_size only when K % block_size == 0.

For contrast, the CPU kernel rounds up per row and decomposes the index before dividing — onnxruntime/contrib_ops/cpu/quantization/gather_block_quantized.cc:

cpp

auto scale_full_block = (quantize_axis_dim + block_size_ - 1) / block_size_ * quantize_N;  // ceil
...
int64_t scale_idx = x * scale_full_block + y / block_size_ * quantize_N + z;

Two further notes:

This affects 4-bit as well as 8-bit data, so it is reachable from a stock build with the stock MatMulNBitsQuantizer (the repro below uses 4-bit only).
The existing coverage does not catch it. TestOpMatMul4Bits::test_quantize_gather_int4_symmetric / _offsets use embedding_len=228 (not a multiple of 32) but run check_model_correctness at rtol=0.2, atol=0.5, which is far looser than the error this produces on data of magnitude ~0.01.

Real transformer embedding tables are block-aligned, so this does not usually bite in practice — but anything with an unaligned quantized axis is silently wrong on CUDA.

To reproduce

Quantize a Gather to 4 bits with the stock quantizer and compare CPU to CUDA, sweeping the hidden dimension across multiples and non-multiples of block_size:

python

import numpy as np
import onnxruntime as ort
from onnx import TensorProto, helper
from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer

BLOCK = 32
VOCAB = 64


def build_dense(hidden):
    table = (np.random.default_rng(0).standard_normal((VOCAB, hidden)) * 0.01).astype(np.float32)
    initializer = helper.make_tensor("table", TensorProto.FLOAT, table.shape, table.tobytes(), True)
    node = helper.make_node("Gather", ["table", "ids"], ["out"], name="G", axis=0)
    graph = helper.make_graph(
        [node],
        "g",
        [helper.make_tensor_value_info("ids", TensorProto.INT64, ["s"])],
        [helper.make_tensor_value_info("out", TensorProto.FLOAT, ["s", hidden])],
        [initializer],
    )
    model = helper.make_model(
        graph, opset_imports=[helper.make_opsetid("", 21), helper.make_opsetid("com.microsoft", 1)]
    )
    return model, table


ids = np.arange(VOCAB, dtype=np.int64)
print(f"onnxruntime {ort.__version__}   block_size={BLOCK}, bits=4, gather_axis=0, quantize_axis=1\n")
print(f"{'hidden':>7} {'hidden%32':>10} {'CPU max|err|':>13} {'CUDA max|err|':>14} {'CUDA sign flips':>16}")

for hidden in (224, 228, 256, 260, 512, 520):
    model, table = build_dense(hidden)
    quantizer = MatMulNBitsQuantizer(
        model=model,
        bits=4,
        block_size=BLOCK,
        is_symmetric=True,
        op_types_to_quantize=("MatMul", "Gather"),
        algo_config=None,
    )
    quantizer.process()
    quantized = quantizer.model.model.SerializeToString()

    cells = []
    for provider in ("CPUExecutionProvider", "CUDAExecutionProvider"):
        got = ort.InferenceSession(quantized, providers=[provider]).run(None, {"ids": ids})[0]
        err = np.abs(got - table).max()
        flips = int((((got * table) < 0) & (np.abs(table) > 1e-4)).sum())
        cells.append((err, flips))

    print(f"{hidden:>7} {hidden % BLOCK:>10} {cells[0][0]:>13.5f} {cells[1][0]:>14.5f} {cells[1][1]:>16}")

Output:

onnxruntime 1.31.0   block_size=32, bits=4, gather_axis=0, quantize_axis=1

 hidden  hidden%32  CPU max|err|  CUDA max|err|  CUDA sign flips
    224          0       0.00365        0.00365                0
    228          4       0.00365        0.06903             6704
    256          0       0.00365        0.00365                0
    260          4       0.00365        0.06876             7260
    512          0       0.00365        0.00365                0
    520          8       0.00365        0.07261            14720

Every row with hidden % block_size != 0 is wrong on CUDA and correct on CPU. "Sign flips" counts output elements whose sign differs from the original dense table (ignoring near-zero values), so these are not rounding differences.

Urgency

No blocker. Models with a block-aligned quantized axis — which includes essentially all transformer embedding tables — are unaffected. Filing because it is a silent correctness bug that current tests are too loose to detect.

Platform

Linux

OS Version

Ubuntu 24.04

ONNX Runtime Installation

Built from Source

ONNX Runtime Version or Commit ID

1.31.0

ONNX Runtime API

Python

Architecture

X64

Execution Provider

CUDA

Execution Provider Library Version

CUDA 13.0, driver 580.159.04, NVIDIA H200 (SM90)

Source: microsoft/onnxruntime

View original on GitHub View discussion on GitHub