Fix int8 GEMM by using dedicated BLAS Lt handle on ROCm/CUDA

csrc/ops.cu

@@ -8,8 +8,6 @@
 #include <limits>
 #include <ops.cuh>
 
-#define ERR_NOT_IMPLEMENTED 100
-
 #if BNB_HIP
 #include <atomic>
 #include <hip/hip_runtime.h>
@@ -407,6 +405,17 @@ int fill_up_to_nearest_multiple(int value, int multiple) {
     return value + (value % multiple == 0 ? 0 : (multiple - (value % multiple)));
 }
 
+#if BNB_HIP && defined(NO_HIPBLASLT)
+int igemmlt_32_gemmex_fallback(
+    Context* context, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, int lda, int ldb, int ldc,
+    bnb_stream_t stream
+) {
+    rocblas_set_stream(context->m_handle, stream);
+    gemmex(context, true, false, k, n, m, (void*)A, (void*)B, (void*)C, lda, ldb, ldc);
+    return 0;
+}
+#endif
+
 void dequant_mm_int32_fp16(
     int* A, float* rowStats, float* colStats, half* out, half* bias, int numRows, int numCols, bnb_stream_t stream
 ) {

csrc/ops.cuh

@@ -6,6 +6,8 @@
 #ifndef ops_H
 #define ops_H
 
+#define ERR_NOT_IMPLEMENTED 100
+
 #include <assert.h>
 #include <cstdint>
 #include <functional>
@@ -62,31 +64,30 @@ class Context {
   public:
 #if BNB_HIP
     rocblas_handle m_handle;
+#else
+    cublasHandle_t m_handle;
+#endif
+    // Separate Lt handle for igemmlt (int8 matmul). Do not reuse m_handle: on HIP it is
+    // rocblas/hipblas, on CUDA it is cublas — neither is compatible with (hip|cublas)Lt APIs.
+#if !BNB_HIP || !defined(NO_HIPBLASLT)
+    bnb_blasLt_handle_t m_lt_handle;
+#endif
 
     Context() {
+#if BNB_HIP
         rocblas_handle handle;
         rocblas_create_handle(&handle);
         m_handle = handle;
-    }
 #else
-    cublasHandle_t m_handle;
-
-    Context() {
         cublasHandle_t handle;
         cublasCreate_v2(&handle);
         m_handle = handle;
-    }
 #endif
-};
-
-class ContextLt {
-  public:
-    bnb_blasLt_handle_t m_handle;
-
-    ContextLt() {
-        bnb_blasLt_handle_t handle;
-        bnb_blasLtCreate(&handle);
-        m_handle = handle;
+#if !BNB_HIP || !defined(NO_HIPBLASLT)
+        bnb_blasLt_handle_t lt_handle;
+        bnb_blasLtCreate(&lt_handle);
+        m_lt_handle = lt_handle;
+#endif
     }
 };
 
@@ -130,6 +131,14 @@ int igemmlt(
     int lda, int ldb, int ldc, bnb_stream_t stream
 );
 
+#if BNB_HIP && defined(NO_HIPBLASLT)
+// ROCm builds without hipBLASLt: int8 matmul via hipblasGemmEx on the rocblas handle.
+int igemmlt_32_gemmex_fallback(
+    Context* context, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, int lda, int ldb, int ldc,
+    bnb_stream_t stream
+);
+#endif
+
 void cutlass_igemm(
     bool transposeA, bool transposeB, int m, int n, int k, void* A, void* B, void* C, int lda, int ldb, int ldc
 );

csrc/pythonInterface.cpp

@@ -525,21 +525,33 @@ int cigemmlt_32(
     Context* context, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda,
     int ldb, int ldc, cudaStream_t stream
 ) {
-    return igemmlt_32((cublasLtHandle_t)context->m_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream);
+#if BUILD_HIP && defined(NO_HIPBLASLT)
+    return igemmlt_32_gemmex_fallback(context, m, n, k, A, B, C, lda, ldb, ldc, stream);
+#else
+    return igemmlt_32(context->m_lt_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream);
+#endif
 }
 
 int cigemmlt_8(
     Context* context, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda,
     int ldb, int ldc, cudaStream_t stream
 ) {
-    return igemmlt_8((cublasLtHandle_t)context->m_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream);
+#if BUILD_HIP && defined(NO_HIPBLASLT)
+    return ERR_NOT_IMPLEMENTED;
+#else
+    return igemmlt_8(context->m_lt_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream);
+#endif
 }
 
 int cigemmlt_8_rowscale(
     Context* context, int m, int n, int k, const int8_t* A, const int8_t* B, void* C, float* row_scale, int lda,
     int ldb, int ldc, cudaStream_t stream
 ) {
-    return igemmlt_8_rowscale((cublasLtHandle_t)context->m_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream);
+#if BUILD_HIP && defined(NO_HIPBLASLT)
+    return ERR_NOT_IMPLEMENTED;
+#else
+    return igemmlt_8_rowscale(context->m_lt_handle, m, n, k, A, B, C, row_scale, lda, ldb, ldc, stream);
+#endif
 }
 
 void cdequant_mm_int32_fp16(

tests/test_cuda_setup_evaluator.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from bitsandbytes.cextension import get_cuda_bnb_library_path
+from bitsandbytes.cextension import HIP_ENVIRONMENT, get_cuda_bnb_library_path
 from bitsandbytes.consts import DYNAMIC_LIBRARY_SUFFIX
 from bitsandbytes.cuda_specs import CUDASpecs
 
@@ -106,7 +106,9 @@ def test_version_selection(monkeypatch, caplog, spec, fake_libs, hip_version, ex
 
 
 def test_override(monkeypatch, cuda120_spec, caplog):
-    """BNB_CUDA_VERSION overrides path selection."""
+    """BNB_CUDA_VERSION overrides path selection (CUDA builds only)."""
+    if HIP_ENVIRONMENT:
+        pytest.skip("BNB_CUDA_VERSION applies to CUDA library selection; use test_rocm_override on ROCm")
     monkeypatch.setenv("BNB_CUDA_VERSION", "110")
     with patch("bitsandbytes.cextension._find_cuda_libs", return_value={}):
         with caplog.at_level("WARNING"):
@@ -130,6 +132,21 @@ def test_rocm_override(monkeypatch, rocm70_spec, caplog):
 
 def test_override_invalid_format(monkeypatch, cuda120_spec):
     """Override value must be digits only (e.g. '124'), not dotted or alphanumeric."""
+    if HIP_ENVIRONMENT:
+        pytest.skip("BNB_CUDA_VERSION applies to CUDA library selection; use test_rocm_override_invalid_format on ROCm")
     monkeypatch.setenv("BNB_CUDA_VERSION", "12.4")
     with pytest.raises(RuntimeError, match="digits only"):
         get_cuda_bnb_library_path(cuda120_spec)
+
+
+def test_rocm_override_invalid_format(monkeypatch, rocm70_spec):
+    """BNB_ROCM_VERSION override must be digits only on ROCm."""
+    if not HIP_ENVIRONMENT:
+        pytest.skip("BNB_ROCM_VERSION applies to ROCm library selection only")
+    monkeypatch.setenv("BNB_ROCM_VERSION", "7.2")
+    with (
+        patch("torch.version.hip", "7.0.0"),
+        patch("bitsandbytes.cextension._find_cuda_libs", return_value={}),
+    ):
+        with pytest.raises(RuntimeError, match="digits only"):
+            get_cuda_bnb_library_path(rocm70_spec)

tests/test_ops.py

@@ -20,8 +20,64 @@ def test_int8_linear_matmul(self, device):
         assert out.dtype == torch.int32
         assert out.device == A.device
 
+        ref = torch.matmul(A.float(), B.t().float()).to(torch.int32)
+        torch.testing.assert_close(out, ref)
+
         opcheck(torch.ops.bitsandbytes.int8_linear_matmul.default, (A, B))
 
+    @pytest.mark.parametrize("device", get_available_devices())
+    def test_int8_linear_matmul_preserves_quant_stats(self, device):
+        """Regression: int8 matmul must use the (hip|cublas)Lt handle, not the BLAS handle.
+
+        Passing a rocblas/cublas handle to hipBLASLt/cuBLASLt caused OOB writes that
+        corrupted unrelated float32 quantization stats sitting in GPU memory.
+        """
+        if device == "cpu":
+            pytest.skip("int8 Lt matmul is GPU-only")
+
+        torch.manual_seed(0)
+        rows, inner, out_features = 256, 128, 64
+        activations = torch.randn(rows, inner, device=device)
+        weights = torch.randn(out_features, inner, device=device)
+        row_stats = torch.amax(torch.abs(activations), dim=1, keepdim=True)
+        col_stats = torch.amax(torch.abs(weights), dim=1, keepdim=True)
+        row_stats_before = row_stats.clone()
+        col_stats_before = col_stats.clone()
+
+        activations_q = torch.round(activations * (127.0 / row_stats)).to(torch.int8)
+        weights_q = torch.round(weights * (127.0 / col_stats)).to(torch.int8)
+
+        out = torch.ops.bitsandbytes.int8_linear_matmul.default(activations_q, weights_q)
+        torch.cuda.synchronize()
+
+        assert out.shape == (rows, out_features)
+        assert not torch.isnan(out.float()).any()
+        torch.testing.assert_close(row_stats, row_stats_before)
+        torch.testing.assert_close(col_stats, col_stats_before)
+
+    @pytest.mark.parametrize("device", get_available_devices())
+    def test_int8_scaled_mm_matches_reference(self, device):
+        """int8_scaled_mm composes Lt int8 matmul + dequant; validate end-to-end numerics."""
+        if device == "cpu":
+            pytest.skip("int8 Lt matmul is GPU-only")
+
+        torch.manual_seed(0)
+        rows, inner, out_features = 256, 128, 64
+        activations = torch.randn(rows, inner, device=device, dtype=torch.float16)
+        weights = torch.randn(out_features, inner, device=device, dtype=torch.float16)
+        row_stats = torch.amax(torch.abs(activations.float()), dim=1)
+        col_stats = torch.amax(torch.abs(weights.float()), dim=1)
+        activations_q = torch.round(activations.float() * (127.0 / row_stats.view(-1, 1))).to(torch.int8)
+        weights_q = torch.round(weights.float() * (127.0 / col_stats.view(-1, 1))).to(torch.int8)
+
+        out = torch.ops.bitsandbytes.int8_scaled_mm.default(
+            activations_q, weights_q, row_stats, col_stats, bias=None, dtype=torch.float16
+        )
+        ref = torch.matmul(activations, weights.t())
+
+        assert out.isfinite().all()
+        assert (out - ref).abs().mean() < 0.1
+
     @pytest.mark.parametrize("device", get_available_devices())
     def test_int8_linear_matmul_out(self, device):
         A = torch.randint(-128, 127, (10, 20), dtype=torch.int8, device=device)