From 775df9d2c21976512e5ad98fa12534cb368a65ca Mon Sep 17 00:00:00 2001
From: Wang Zupeng <zupenwang@gmail.com>
Date: Fri, 23 Oct 2026 19:47:39 +0800
Subject: [PATCH] Fix CUDA INT8 quantization of non-contiguous inputs

---
 bitsandbytes/backends/cuda/ops.py |  3 +++
 tests/test_linear8bitlt.py        | 26 ++++++++++++++++++++++
 tests/test_ops.py                 | 37 +++++++++++++++++++++++++++++++
 3 files changed, 66 insertions(+)

diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py
index a0d9ffe83..ad87b30b6 100644
--- a/bitsandbytes/backends/cuda/ops.py
+++ b/bitsandbytes/backends/cuda/ops.py
@@ -217,6 +217,9 @@ def _(A: torch.Tensor, threshold=0.0):
     if threshold < 0.0:
         raise ValueError("threshold must be non-negative")
 
+    # The CUDA kernel reads rows directly without accounting for tensor strides.
+    A = A.contiguous()
+
     rows = A.numel() // A.shape[-1]
     cols = A.shape[-1]
 
diff --git a/tests/test_linear8bitlt.py b/tests/test_linear8bitlt.py
index b078e82b7..404b27b02 100644
--- a/tests/test_linear8bitlt.py
+++ b/tests/test_linear8bitlt.py
@@ -22,6 +22,32 @@
 HAS_FP16_WEIGHTS_NON_DEPRECATED = [False]
 
 
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("threshold", [0.0, 6.0])
+@pytest.mark.parametrize("layout", ["transpose", "column_slice"])
+@pytest.mark.parametrize("requires_grad", TRUE_FALSE)
+def test_linear8bit_non_contiguous_input(device, threshold, layout, requires_grad):
+    linear = Linear8bitLt(64, 32, has_fp16_weights=False, threshold=threshold).half().to(device).eval()
+    if layout == "transpose":
+        x = torch.randn(64, 4, device=device, dtype=torch.float16).t()
+    else:
+        x = torch.randn(4, 128, device=device, dtype=torch.float16)[:, 1::2]
+    x[1, 3] = 10.0
+    assert not x.is_contiguous()
+    x_ref = x.contiguous().requires_grad_(requires_grad)
+    x.requires_grad_(requires_grad)
+
+    expected = linear(x_ref)
+    actual = linear(x)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+    if requires_grad:
+        grad = torch.randn_like(actual)
+        expected.backward(grad)
+        actual.backward(grad)
+        torch.testing.assert_close(x.grad, x_ref.grad, rtol=0, atol=0)
+
+
 # contributed by Alex Borzunov, see:
 # https://github.com/bigscience-workshop/petals/blob/main/tests/test_linear8bitlt.py
 @pytest.mark.parametrize("device", get_available_devices())
diff --git a/tests/test_ops.py b/tests/test_ops.py
index 4ca60f845..9e4752118 100644
--- a/tests/test_ops.py
+++ b/tests/test_ops.py
@@ -62,6 +62,43 @@ def test_int8_vectorwise_quant(self, threshold, device):
         opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A,))
         opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold))
 
+    @pytest.mark.parametrize("device", get_available_devices())
+    @pytest.mark.parametrize("threshold", [0.0, 6.0])
+    @pytest.mark.parametrize("layout", ["transpose", "row_slice", "column_slice"])
+    @pytest.mark.parametrize("double_quant", TRUE_FALSE)
+    def test_int8_quant_non_contiguous(self, device, threshold, layout, double_quant):
+        if double_quant and device != "cuda":
+            pytest.skip("int8_double_quant is only implemented for CUDA/ROCm")
+
+        if layout == "transpose":
+            A = torch.randn(64, 8, device=device, dtype=torch.float16).t()
+        elif layout == "row_slice":
+            A = torch.randn(16, 64, device=device, dtype=torch.float16)[1::2]
+        else:
+            A = torch.randn(8, 128, device=device, dtype=torch.float16)[:, 1::2]
+        A[1, 3] = 10.0
+        A[6, 17] = -10.0
+        assert not A.is_contiguous()
+        original = A.clone()
+
+        quantize = (
+            bitsandbytes.functional.int8_double_quant
+            if double_quant
+            else bitsandbytes.functional.int8_vectorwise_quant
+        )
+        expected = quantize(A.contiguous(), threshold=threshold)
+        actual = quantize(A, threshold=threshold)
+
+        for result, reference in zip(actual, expected):
+            if reference is None:
+                assert result is None
+            else:
+                torch.testing.assert_close(result, reference, rtol=0, atol=0)
+        torch.testing.assert_close(A, original, rtol=0, atol=0)
+
+        if device == "cuda" and not double_quant:
+            opcheck(torch.ops.bitsandbytes.int8_vectorwise_quant, (A, threshold))
+
     @pytest.mark.parametrize("device", get_available_devices())
     def test_int8_mm_dequant(self, device):
         A = torch.randint(-128, 127, (256, 256), dtype=torch.int32, device=device)