From 496facb5728f0d2a139095bebeb79c195533d62e Mon Sep 17 00:00:00 2001 From: Ege Ozkoc <90700230+egeozkoc@users.noreply.github.com> Date: Sun, 8 Nov 2026 17:27:56 -0500 Subject: [PATCH] test(mps): verify 4-bit dequantization values and Hub dispatch --- tests/conftest.py | 60 +++++++++++++++++++++++++++++++++++ tests/test_ops.py | 79 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 139 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index 0d2e162f3..d90b7c2f7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,66 @@ import torch +def pytest_addoption(parser): + parser.addoption( + "--require-mps-hub", + action="store_true", + help="Require MPS, the bitsandbytes-mps Hub kernel, and a passing test that verifies Hub execution.", + ) + + +def _require_mps_hub(): + from tests.helpers import get_available_devices + + if "mps" not in get_available_devices() or not torch.backends.mps.is_available(): + raise pytest.UsageError("--require-mps-hub requires an available MPS device selected by BNB_TEST_DEVICE.") + + from bitsandbytes.backends.mps import ops + + kernel = ops._get_kernel() + if kernel is None: + raise pytest.UsageError( + "--require-mps-hub could not load kernels-community/bitsandbytes-mps; " + "macOS 26+ and a compatible kernels package/Hub build are required." + ) + return kernel + + +class _MPSHubExecutionCheck: + def __init__(self): + self.verified_tests = 0 + + def pytest_runtest_logreport(self, report): + if report.when == "call" and report.passed and ("mps_hub_verified", True) in report.user_properties: + self.verified_tests += 1 + + def pytest_sessionfinish(self, session, exitstatus): + if exitstatus == pytest.ExitCode.OK and not self.verified_tests: + session.exitstatus = pytest.ExitCode.TESTS_FAILED + reporter = session.config.pluginmanager.get_plugin("terminalreporter") + if reporter is not None: + reporter.write_sep( + "=", + "ERROR: --require-mps-hub completed without any passing test verifying Hub execution. " + "Include test_dequantize_4bit_mps with a supported Hub block size.", + red=True, + ) + + +def pytest_configure(config): + if config.getoption("--require-mps-hub"): + _require_mps_hub() + config.pluginmanager.register(_MPSHubExecutionCheck(), "mps-hub-execution-check") + + +def pytest_report_header(config): + if config.getoption("--require-mps-hub"): + from importlib.metadata import version + + kernel = _require_mps_hub() + return f"MPS Hub required: torch={torch.__version__}, kernels={version('kernels')}, variant={kernel.__file__}" + + def _set_seed(): torch.manual_seed(0) torch.cuda.manual_seed_all(0) diff --git a/tests/test_ops.py b/tests/test_ops.py index 4ca60f845..94fe898d8 100644 --- a/tests/test_ops.py +++ b/tests/test_ops.py @@ -1,4 +1,5 @@ from math import prod +from unittest.mock import Mock import pytest import torch @@ -9,6 +10,42 @@ opcheck = torch.library.opcheck +def _dequantize_4bit_reference(packed, absmax, blocksize, quant_type, shape, dtype): + # Decode on CPU independently of the quantizer and all backend dequantizers. + raw = packed.cpu().contiguous().view(torch.uint8).reshape(-1) + codes = torch.stack((raw >> 4, raw & 0xF), dim=1).flatten()[: prod(shape)].long() + codebook = bitsandbytes.functional.get_4bit_type(quant_type, device="cpu") + scales = absmax.cpu().repeat_interleave(blocksize)[: prod(shape)] + return (codebook[codes] * scales).reshape(shape).to(dtype) + + +@pytest.fixture(params=["hub", "fallback"]) +def mps_dequantize_backend(request, monkeypatch): + if "mps" not in get_available_devices() or not torch.backends.mps.is_available(): + pytest.skip("MPS is not available or not selected for this run") + + from bitsandbytes.backends.mps import ops + + # Restore loader state too: _get_kernel() can mutate both globals. + monkeypatch.setattr(ops, "_kernel", ops._kernel) + monkeypatch.setattr(ops, "_kernel_load_failed", ops._kernel_load_failed) + kernel = ops._get_kernel() + if request.param == "hub" and kernel is None: + if request.config.getoption("--require-mps-hub"): + pytest.fail("The required MPS Hub kernel is unavailable") + pytest.skip("The bitsandbytes-mps Hub kernel is unavailable (requires macOS 26+ and kernels)") + + if kernel is not None: + calls = Mock(wraps=kernel.dequantize_4bit) + monkeypatch.setattr(kernel, "dequantize_4bit", calls) + else: + calls = Mock(side_effect=AssertionError("Hub dequantization must not run in fallback mode")) + monkeypatch.setattr(ops, "_kernel", Mock(dequantize_4bit=calls)) + if request.param == "fallback": + monkeypatch.setattr(ops, "_kernel_load_failed", True) + return request.param, calls + + class TestLLMInt8Ops: @pytest.mark.parametrize("device", get_available_devices()) def test_int8_linear_matmul(self, device): @@ -213,12 +250,54 @@ def test_dequantize_4bit(self, device, dtype, storage_dtype, quant_type, blocksi assert out.device == A.device assert out.shape == shape + assert out.dtype == dtype + expected = _dequantize_4bit_reference(A, absmax, blocksize, quant_type, shape, dtype) + torch.testing.assert_close(out.cpu(), expected) opcheck( torch.ops.bitsandbytes.dequantize_4bit.default, (A, absmax, blocksize, quant_type, shape, dtype), ) + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) + @pytest.mark.parametrize("quant_type", ["fp4", "nf4"]) + @pytest.mark.parametrize( + ("blocksize", "shape", "storage_dtype"), + [ + (64, (4, 64), torch.uint8), + (128, (4, 128), torch.bfloat16), + (256, (3, 257), torch.bfloat16), + (512, (3, 513), torch.uint8), + (32, (3, 33), torch.bfloat16), + ], + ids=["block64-aligned", "block128-aligned", "block256-tail", "block512-tail", "block32-fallback"], + ) + def test_dequantize_4bit_mps( + self, mps_dequantize_backend, record_property, dtype, quant_type, blocksize, shape, storage_dtype + ): + mode, hub_calls = mps_dequantize_backend + hub_calls.reset_mock() + indices = torch.arange((prod(shape) + 1) // 2) + # Exercise every nibble code, including different high/low nibbles. + raw = (((indices % 16) << 4) | ((indices * 7 + 3) % 16)).to(torch.uint8) + packed = raw.view(storage_dtype).reshape(-1, 1).to("mps") + absmax = torch.tensor([0.0, 0.37, 1.125, 2.75], device="mps") + expected = _dequantize_4bit_reference(packed, absmax, blocksize, quant_type, shape, dtype) + + actual = torch.ops.bitsandbytes.dequantize_4bit.default(packed, absmax, blocksize, quant_type, shape, dtype) + assert actual.device == packed.device + torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0) + + out = torch.empty((shape[1], shape[0]), device="mps", dtype=dtype).T + pointer = out.data_ptr() + torch.ops.bitsandbytes.dequantize_4bit.out(packed, absmax, blocksize, quant_type, shape, dtype, out=out) + assert out.data_ptr() == pointer + assert not out.is_contiguous() + torch.testing.assert_close(out.cpu(), expected, rtol=0, atol=0) + assert hub_calls.call_count == (2 if mode == "hub" and blocksize != 32 else 0) + if mode == "hub" and blocksize != 32: + record_property("mps_hub_verified", True) + @pytest.mark.parametrize("device", get_available_devices()) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32], ids=id_formatter("dtype")) @pytest.mark.parametrize("storage_dtype", [torch.uint8, torch.bfloat16], ids=id_formatter("storage_dtype"))