Skip to content

Fix CUDA INT8 quantization of non-contiguous inputs - #2104

Open
zupengwang wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
zupengwang:fix-int8-quantize-strided-input
Open

zupengwang wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
zupengwang:fix-int8-quantize-strided-input

Conversation

@zupengwang

Copy link
Copy Markdown

The CUDA implementation of int8_vectorwise_quant passes the input pointer directly to a kernel that assumes contiguous rows. Transposed inputs and row/column slices therefore produce incorrect quantized values and row scales without raising an error. This also affects int8_double_quant and FP16 inputs to Linear8bitLt.

Make the input contiguous before calling the CUDA kernel. Add regression coverage for transposes, strided slices with nonzero storage offsets, thresholds 0 and 6, both quantization functions, and Linear8bitLt outputs and input gradients. The CUDA vectorwise tests also run torch.library.opcheck.

Minimal reproducer on the unpatched CUDA backend:

import torch
from bitsandbytes.functional import int8_vectorwise_quant

torch.manual_seed(0)
x = torch.randn(64, 8, device="cuda", dtype=torch.float16).t()
actual, scales, _ = int8_vectorwise_quant(x)
expected, expected_scales, _ = int8_vectorwise_quant(x.contiguous())
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
torch.testing.assert_close(scales, expected_scales, rtol=0, atol=0)

Validation against main 833649043474794b8fe7a4136e0c40faf077b2e0, using a source build on RTX 4090 (SM89), Python 3.12.14, PyTorch 2.13.0+cu130, and CUDA 13.0:

  • All 20 new CUDA regression cases fail before the fix and pass after it, using the same native library.
  • BNB_TEST_DEVICE=cuda python -m pytest tests/test_ops.py::TestLLMInt8Ops tests/test_linear8bitlt.py -q -ra: 84 passed, 2 existing xfails for torch.compile reduce-overhead mode. This total includes the 20 new cases.
  • The two new tests with BNB_TEST_DEVICE=cpu: 14 passed, 6 skipped because double quantization is CUDA/ROCm-only. opcheck in the new test is scoped to CUDA; the default CPU backend has a separate existing output-stride mismatch with its fake implementation for transposed inputs.
  • Ruff 0.14.3 lint/format checks and git diff --check passed.

These are local hardware results; upstream CI and other accelerators have not been run.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant