Cpu C++ kernel

[jiqing-feng]

The C++ kernels.

cmake -DCOMPUTE_BACKEND=cpu -S . && make

Hi @matthewdouglas . I've implemented the CPU dequantize op for nf4/fp4. It will bring 10x+ speed-up in the e2e text-generation task compared with the original python kernel on llama3-8B model. Would you please review this PR? Thanks!

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[jiqing-feng]

Hi @matthewdouglas . Please trigger the tests and review the PR. Thanks!

[jiqing-feng]

Hi @matthewdouglas . As we discussed, the biggest problem is that C++ kernels can be build and ran in different platforms. To solve this issue, I added runtime check for avx512f and avx512bf16. It can guarantee the kernel will go to fallback path in the CPU without avx512f. Please review this change and let me know if you have any other concerns.

[jiqing-feng]

Hi @matthewdouglas . I've enabled CPU C++ kernels on Windows, and the following script works well on both linux and Windows

from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
import torch

quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
        )

model_id = "JackFram/llama-68m"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, device_map="cpu", quantization_config=quantization_config, torch_dtype=torch.bfloat16,
)

generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="cpu")

prompt = "Once upon a time in a small village,"
outputs = generator(prompt, max_new_tokens=10, temperature=0.8, top_p=0.95, do_sample=True)

print("=== Generated Text ===")
print(outputs[0]["generated_text"])

building command on linux: cmake -DCOMPUTE_BACKEND=cpu -S . && make
building command on Windows: cmake -DCOMPUTE_BACKEND=cpu -S . && cmake --build . --config Release

[matthewdouglas]

Thanks for the effort here @jiqing-feng!

I've built and tested on Linux with my Ryzen 7950X (supports AVX512-BF16). Most tests are passing, but there's some failures:

This set of failures is for higher blocksizes (2048, 4096). Something to note is that fp32/fp16 have high errors outside the test tolerances, while bf16 passes. It would be good if we can get these back within the existing tolerances.

tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[2048-fp4-fp32-cpu] FAILED [ 58%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[2048-fp4-fp16-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[2048-nf4-fp32-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[2048-nf4-fp16-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[4096-fp4-fp32-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[4096-fp4-fp16-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[4096-nf4-fp32-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[4096-nf4-fp16-cpu] FAILED [ 59%]

These failures seem to be because of a shape with odd number like ([128,65]) and/or 3D shapes and should be fixed:

tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.uint8-(10,)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.uint8-(10, 10)-64-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.uint8-(10, 10)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.uint8-(10, 10, 10)-64-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.uint8-(10, 10, 10)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.float32-(10,)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.float32-(10, 10)-64-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.float32-(10, 10)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.float32-(10, 10, 10)-64-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingFP4-torch.float32-(10, 10, 10)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.uint8-(10,)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.uint8-(10, 10)-64-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.uint8-(10, 10)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.uint8-(10, 10, 10)-64-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.uint8-(10, 10, 10)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.float32-(10, 10)-64-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.float32-(10, 10)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.float32-(10, 10, 10)-64-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_lossless[EmbeddingNF4-torch.float32-(10, 10, 10)-65-cpu] FAILED [ 87%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.uint8-(10,)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.uint8-(10, 10)-64-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.uint8-(10, 10)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.uint8-(10, 10, 10)-64-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.uint8-(10, 10, 10)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.float32-(10,)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.float32-(10, 10)-64-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.float32-(10, 10)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.float32-(10, 10, 10)-64-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingFP4-torch.float32-(10, 10, 10)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.uint8-(10,)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.uint8-(10, 10)-64-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.uint8-(10, 10)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.uint8-(10, 10, 10)-64-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.uint8-(10, 10, 10)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.float32-(10,)-64-cpu] PASSED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.float32-(10,)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.float32-(10, 10)-64-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.float32-(10, 10)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.float32-(10, 10, 10)-64-cpu] FAILED [ 88%]
tests/test_modules.py::test_embedding_error[EmbeddingNF4-torch.float32-(10, 10, 10)-65-cpu] FAILED [ 88%]
tests/test_modules.py::test_4bit_embedding_warnings[cpu] FAILED          [ 88%]

This last set of failures below is related to newer functionality and shouldn't block merging; can fix this separately. With that said some of these are probably failing for the same reasons as the embedding ones.

tests/test_parametrize.py::test_moe_parameter_shape[fp32-cpu] FAILED     [ 98%]
tests/test_parametrize.py::test_moe_parameter_shape[fp16-cpu] FAILED     [ 98%]
tests/test_parametrize.py::test_moe_parameter_shape[bf16-cpu] FAILED     [ 98%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=T-nf4-fp32-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=T-nf4-fp16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=T-nf4-bf16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=T-fp4-fp32-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=T-fp4-fp16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=T-fp4-bf16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=F-nf4-fp32-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=F-nf4-fp16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=F-nf4-bf16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_multiple_parameters[fp32-cpu] FAILED     [ 99%]
tests/test_parametrize.py::test_multiple_parameters[fp16-cpu] FAILED     [ 99%]
tests/test_parametrize.py::test_multiple_parameters[bf16-cpu] FAILED     [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=F-fp4-fp32-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=F-fp4-fp16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_state_dict_functionality[compress_statistics=F-fp4-bf16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_moe_realistic_forward[fp32-cpu] FAILED   [ 99%]
tests/test_parametrize.py::test_moe_realistic_forward[fp16-cpu] FAILED   [ 99%]
tests/test_parametrize.py::test_moe_realistic_forward[bf16-cpu] FAILED   [ 99%]
tests/test_parametrize.py::test_different_blocksizes[64-fp32-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_different_blocksizes[64-fp16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_different_blocksizes[64-bf16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_different_blocksizes[128-fp32-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_different_blocksizes[128-fp16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_different_blocksizes[128-bf16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_different_blocksizes[256-fp32-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_different_blocksizes[256-fp16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_different_blocksizes[256-bf16-cpu] FAILED [ 99%]
tests/test_parametrize.py::test_parametrization_forward_method FAILED    [ 99%]

[matthewdouglas]

@jiqing-feng I've pushed a change to take care of most of the test failures. There's 8 failures remaining.

The ones remaining pass without AVX512, but fail with it:

tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[2048-fp4-fp32-cpu] FAILED [ 58%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[2048-fp4-fp16-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[2048-nf4-fp32-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[2048-nf4-fp16-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[4096-fp4-fp32-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[4096-fp4-fp16-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[4096-nf4-fp32-cpu] FAILED [ 59%]
tests/test_functional.py::TestQuantize4BitFunctional::test_4bit_quant[4096-nf4-fp16-cpu] FAILED [ 59%]

I believe these all pass when we use larger tensors than the [1024,1024] used in the test. Moreover, these blocksizes aren't typically used e.g. in transformers integration. I think I'm happy to merge, unless you see this as an issue?

[jiqing-feng]

Hi @matthewdouglas . These failed tests are not related to my changes. We can merge this PR and fix these tests in a seperate PR.

[matthewdouglas]

@jiqing-feng Those tests passed for me when I built without AVX512. They also passed on the nightly CI on main:
https://github.com/bitsandbytes-foundation/bitsandbytes/actions/runs/19318858303

However, I don't think it's a significant blocker, so I'm going to merge this and we'll fix outside of this PR.