Skip to content

test(mps): verify 4-bit dequantization values and Hub dispatch - #2123

Open
egeozkoc wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
egeozkoc:test/mps-dequantization-coverage
Open

egeozkoc wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
egeozkoc:test/mps-dequantization-coverage

Conversation

@egeozkoc

@egeozkoc egeozkoc commented Nov 8, 2026

Copy link
Copy Markdown
Contributor

The direct dequantize_4bit operator test checks shape/device and runs opcheck,
but can pass when the MPS Hub kernel returns zeros. It can also pass entirely
through the fallback, so a successful MPS run does not establish Hub coverage.

Add an independent CPU reference that unpacks high/low nibbles, indexes the
public FP4/NF4 codebook, applies block scales in FP32, and casts to the requested
dtype. The existing operator matrix now checks numerical values and output dtype
before opcheck.

A separate 60-case MPS matrix checks real Hub calls and forced fallback across
FP4/NF4 and FP16/BF16/FP32. It includes all four supported Hub block sizes,
uint8/bfloat16 packed storage, zero/non-unit scales, odd partial tails, and both
allocating and .out calls with a non-contiguous destination. Blocksize 32 must
use fallback even with Hub available. Loader state and the spied function are
restored after each case.

Optional Hub cases skip explicitly if unavailable. --require-mps-hub rejects
missing MPS, an incompatible device selection, or an unavailable Hub build, and
prints Torch/kernels versions and the resolved variant in the pytest header.
The dispatch matrix asserts actual calls, so silent fallback fails its Hub rows.
Required mode also needs at least one passing test that verified Hub execution:
the test records this only after its numerical and dispatch assertions pass,
and a session check rejects an otherwise successful run with none. Selecting
only fallback controls therefore cannot satisfy --require-mps-hub.
Only tests/test_ops.py and tests/conftest.py change.

Validation

Apple M5, macOS 26.0, Python 3.12.13, kernels 0.17.1; current main a616886.
Used the published Hub snapshot c3cadfafa7f693bbabbe65479be68f23ffb1c880, with
its respective Torch 2.12 and 2.15 Metal builds.

On each of PyTorch 2.12.1 and 2.15.0:

  • BNB_TEST_DEVICE=mps python -m pytest tests/test_ops.py -k dequantize_4bit --require-mps-hub: 120 passed.
  • BNB_TEST_DEVICE=cpu python -m pytest tests/test_ops.py -k dequantize_4bit: 60 passed, 60 MPS-only cases skipped, using a freshly built native CPU library.
  • Local mutation controls: returning zeros and swapping packed nibbles each fail
    the numerical assertions in both the existing test and the new Hub test.
    Disabling Hub dispatch after successful setup fails the call assertion.
  • Required-mode controls reject unavailable Hub, unavailable MPS, and CPU-only
    selection. With the kernels import disabled in optional mode, 30 fallback
    cases pass and 30 Hub cases skip
    .
  • Required mode with only the blocksize-32 controls: 12 cases pass, zero Hub
    calls, session exits 1
    with an explicit missing-execution error. Ten
    mutation/availability/selection controls per Torch version behave as expected.

Applicable pre-commit hooks and git diff --check pass. Mutation scripts and raw
validation logs are local evidence outside the proposed PR. The full suite and
CUDA/ROCm/XPU/HPU were not run. Automated upstream MPS execution still needs an
agreed runner; this patch supplies an explicit local/runner validation command.

AI assistance: implementation, local validation, and this description were
prepared with OpenAI Codex.

@egeozkoc

egeozkoc commented Nov 8, 2026

Copy link
Copy Markdown
Contributor Author

@matthewdouglas, could you approve the pending workflows and review this when you have a chance? Local MPS and CPU validation passes on Torch 2.12.1 and 2.15.0, including checks that required-Hub mode rejects fallback-only runs. Details are in the PR description.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant