microsoft / microsoft/onnxruntime
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
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### 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)
Contributor guide
Assessment
This issue has not been assessed yet.