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

Open
#32,663 0 comments 1 reaction 2 assignees Claimed by @tianleiwu View on GitHub
ep:CUDA quantization
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.