Handle non-contiguous inputs in LLM.int8() CUDA ops

bitsandbytes/backends/cuda/ops.py

@@ -87,7 +87,7 @@ def _(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):
 
 
 def _int8_linear_matmul_impl(A: torch.Tensor, B: torch.Tensor, out: torch.Tensor):
-    A, B = B, A
+    A, B = B.contiguous(), A.contiguous()
 
     shapeA = A.shape
     shapeB = B.shape
@@ -183,6 +183,12 @@ def _(
     if col_stats.dtype != torch.float32:
         raise ValueError(f"col_stats must be float32, got {col_stats.dtype}")
 
+    A = A.contiguous()
+    row_stats = row_stats.contiguous()
+    col_stats = col_stats.contiguous()
+    if bias is not None:
+        bias = bias.contiguous()
+
     # Note: cuda kernel only currently supports fp16 output.
     # We'll later cast to desired dtype if needed.
     out = torch.empty_like(A, dtype=torch.float16)
@@ -217,6 +223,8 @@ def _(A: torch.Tensor, threshold=0.0):
     if threshold < 0.0:
         raise ValueError("threshold must be non-negative")
 
+    A = A.contiguous()
+
     rows = A.numel() // A.shape[-1]
     cols = A.shape[-1]
 

tests/test_linear8bitlt.py

@@ -364,3 +364,22 @@ def test_linear8bitlt_device_movement(device):
 
     # Accelerator outputs should match both times.
     torch.testing.assert_close(out_accelerator_2, out_accelerator, rtol=1e-8, atol=1e-8)
+
+
+@pytest.mark.parametrize("device", get_available_devices())
+@pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold"))
+def test_linear8bitlt_non_contiguous_input(device, threshold):
+    # Pooling and last-token heads pass strided views of the hidden states to the layer,
+    # e.g. hidden_states[:, 0] (BERT-style pooler) or hidden_states[:, -1:, :].
+    torch.manual_seed(0)
+    layer = Linear8bitLt(64, 32, has_fp16_weights=False, threshold=threshold).to(device)
+    hidden_states = torch.randn(8, 16, 64, dtype=torch.float16, device=device)
+
+    for x in (hidden_states[:, 0], hidden_states[:, -1:, :]):
+        assert not x.is_contiguous()
+        with torch.no_grad():
+            out = layer(x)
+            ref = layer(x.contiguous())
+        # The bug reads unrelated rows, so the error is O(1); the tolerance only absorbs
+        # last-bit differences between strided and contiguous paths on some backends.
+        torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2)

tests/test_ops.py

@@ -472,3 +472,50 @@ def test_quantize_4bit_roundtrip_non_contiguous(self, device, dtype, quant_type,
         deq_c = torch.ops.bitsandbytes.dequantize_4bit(q_c, absmax_c, blocksize, quant_type, shape, dtype)
 
         torch.testing.assert_close(deq_nc, deq_c)
+
+    @pytest.mark.parametrize("device", get_available_devices())
+    @pytest.mark.parametrize("threshold", [0.0, 6.0], ids=id_formatter("threshold"))
+    def test_int8_vectorwise_quant_non_contiguous(self, device, threshold):
+        # Token pooling (hidden_states[:, 0]) gives a 2D view whose rows are not adjacent in memory.
+        A_full = torch.randn(8, 16, 64, dtype=torch.float16, device=device)
+        A_full[:, :, 3] = 10.0  # an outlier column, so the threshold path is exercised
+        A_noncontig = A_full[:, 0]
+        assert not A_noncontig.is_contiguous()
+
+        A_contig = A_noncontig.contiguous()
+
+        out_nc, stats_nc, outliers_nc = torch.ops.bitsandbytes.int8_vectorwise_quant(A_noncontig, threshold)
+        out_c, stats_c, outliers_c = torch.ops.bitsandbytes.int8_vectorwise_quant(A_contig, threshold)
+
+        torch.testing.assert_close(stats_nc, stats_c)
+        torch.testing.assert_close(out_nc, out_c)
+        if threshold > 0.0:
+            torch.testing.assert_close(outliers_nc, outliers_c)
+
+    @pytest.mark.parametrize("device", get_available_devices())
+    def test_int8_linear_matmul_non_contiguous(self, device):
+        A_noncontig = torch.randint(-128, 127, (10, 3, 32), dtype=torch.int8, device=device)[:, 0]
+        B_noncontig = torch.randint(-128, 127, (32, 30), dtype=torch.int8, device=device).t()
+        assert not A_noncontig.is_contiguous()
+        assert not B_noncontig.is_contiguous()
+
+        out_nc = torch.ops.bitsandbytes.int8_linear_matmul(A_noncontig, B_noncontig)
+        out_c = torch.ops.bitsandbytes.int8_linear_matmul(A_noncontig.contiguous(), B_noncontig.contiguous())
+
+        torch.testing.assert_close(out_nc, out_c)
+
+    @pytest.mark.parametrize("device", get_available_devices())
+    def test_int8_mm_dequant_non_contiguous(self, device):
+        A_noncontig = torch.randint(-128, 127, (4, 32, 64), dtype=torch.int32, device=device)[:, 0]
+        row_stats_noncontig = torch.rand(8, dtype=torch.float32, device=device)[::2]
+        col_stats_noncontig = torch.rand(128, dtype=torch.float32, device=device)[::2]
+        assert not A_noncontig.is_contiguous()
+        assert not row_stats_noncontig.is_contiguous()
+        assert not col_stats_noncontig.is_contiguous()
+
+        out_nc = torch.ops.bitsandbytes.int8_mm_dequant(A_noncontig, row_stats_noncontig, col_stats_noncontig)
+        out_c = torch.ops.bitsandbytes.int8_mm_dequant(
+            A_noncontig.contiguous(), row_stats_noncontig.contiguous(), col_stats_noncontig.contiguous()
+        )
+
+        torch.testing.assert_close(out_nc, out_c)