diff --git a/csrc/ops.cu b/csrc/ops.cu index 16eed4e81..ac8f47fa9 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -8,8 +8,6 @@ #include #include -#define ERR_NOT_IMPLEMENTED 100 - #if BNB_HIP #include #include @@ -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 ) { diff --git a/csrc/ops.cuh b/csrc/ops.cuh index c7114bcaa..2d4e45ad9 100644 --- a/csrc/ops.cuh +++ b/csrc/ops.cuh @@ -6,6 +6,8 @@ #ifndef ops_H #define ops_H +#define ERR_NOT_IMPLEMENTED 100 + #include #include #include @@ -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(<_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 ); diff --git a/csrc/pythonInterface.cpp b/csrc/pythonInterface.cpp index 24431603b..a88202d42 100644 --- a/csrc/pythonInterface.cpp +++ b/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( diff --git a/tests/test_cuda_setup_evaluator.py b/tests/test_cuda_setup_evaluator.py index 56a52736e..6787a8064 100644 --- a/tests/test_cuda_setup_evaluator.py +++ b/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) diff --git a/tests/test_ops.py b/tests/test_ops.py index 4ca60f845..6bf3a226a 100644 --- a/tests/test_ops.py +++ b/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)