Small but mighty!

There is probably space for further optimizations. I have seen slightly faster inference for this model.

But I believe it's a good start!

Implemented kernel caching and reduced precision from float32 to float16, achieving significant performance gains across all metrics:

Key Improvements:

• Throughput: 2x increase overall
• Peak TFLOPS: 45.86 TFLOPS (up from 20.19) — 127% improvement
• MLP Forward Pass: 40% faster execution
• Prompt Processing: 135.046 tokens/sec (up from 51) — 165% improvement
• Generation Speed: 46.209 tokens/sec (up from 43.830) — 5% improvement
• Peak Memory: 1.308 GB (down from 1.322 GB) — 1% reduction

Fusing qkv can probably get us to around ~50 tokens but I couldn't find a way to join the scales and keep model coherence.

Hi @Blaizzy - the results look impressive, this could potentially be easily ported to other existing Bitnet models such as Falcon-E or Falcon3-1.58 since we can use the same BitNetLinear kernel and are based on Llama architecture. I am happy to work on integrating these models as well, would the BitLinear class be re-usable by other models? How do you think we should approach this?

Currently the way to infer if the models use bitnet for these models is to check this attribute: https://huggingface.co/tiiuae/Falcon-E-3B-Instruct/blob/main/config.json#L27

Hi @Blaizzy - the results look impressive, this could potentially be easily ported to other existing Bitnet models such as Falcon-E or Falcon3-1.58 since we can use the same BitNetLinear kernel and are based on Llama architecture. I am happy to work on integrating these models as well, would the BitLinear class be re-usable by other models? How do you think we should approach this?

Thank you @younesbelkada!

Yes, I can put all the BitLinear Layer logic into a separate file that can be re-used. Much like the existing switch_layers.py

Funny you meantion falcon H1, I worked on a MLX version a few weeks back but in a separate repo. I would love to collaborate with you on it.

Currently the way to infer if the models use bitnet for these models is to check this attribute: https://huggingface.co/tiiuae/Falcon-E-3B-Instruct/blob/main/config.json#L27

Awesome, that makes things easier!

@Blaizzy - great to hear that ! I'll reach out to you separately by email (the one you put on your GH profile), I would love to colab as well for H1 and Bitnet ! Sending you an email now

@Blaizzy - great to hear that ! I'll reach out to you separately by email (the one you put on your GH profile), I would love to colab as well for H1 and Bitnet ! Sending you an email now

Perfect!

I have moved BitLinear layer logic to a separate file for easier reusability across models and projects.

Thank you @Blaizzy will give it a try and I'll either open a PR to your branch or on main if this PR gets merged quickly

My pleasure!

You can send a PR to my branch 👌🏽

Got a working version: Blaizzy#1

Also works on previous Bitnet models, e.g. https://huggingface.co/tiiuae/Falcon3-7B-Instruct-1.58bit :

Perfect, great job @younesbelkada ! 🚀

I was about to text you saying I tried Falcon 3 but the hidden states where exploding to Infinity 😂

But the inverse change fixes it ✅

Left a comment in your PR @younesbelkada after that is resolved we can merge.

Thank you @Blaizzy ! Also sent you an email for collab 🙏
Looking forward to seeing this merged 🚀

Performance of the 1B model:

Just added support for N-bit quants for 1.58bit model and it further reduces peak memory and is faster.

@awni @angeloskath

This last commit is small but makes a huge difference!

MLX is now officially faster than bitnet.cpp, without even using N-bit quants that add +10 tokens/s.

And I believe with fused qkv kernels we can get an additional 5-10 tokens/s :)

2T = 2 threads
4T = 4 threads

@Blaizzy your code looks really nice. One question: your 0,1,2 to -1,0,1 shift is one additional operation can we avoid it ? second question: do you plan to include the fused qkv kernels too ?

I'm afraid we can't avoid it.

•	2-bit packing – 4 weights per byte means each weight is an unsigned 2-bit code (00–11).
•	Ternary encoding – store {-1, 0, +1} as {00, 01, 10} and recover with bits - 1.
•	Mat-mul runs in FP16/FP32 – inputs are floats and Apple GPUs multiply-accumulate only in FP16/32 (or INT8 after unpacking), so every weight must be promoted to FP anyway.
•	Negligible cost – float(bits) - 1 is a single fused ALU op; any alternate encoding adds more instructions or an extra reduction.

Yes, I plan to include the fused QKV kernels as soon as I find the best way to aggregate the weight scales for QKV.

Fused QKV kernel is coming soon 🚀

so on M3 128 GB: generate at 89.3 +/- 0.1 tok/sec.
Why the first call prompt is always slower (model loading maybe) ?
I have Prompt first time 52.5 and then 162.2 +/-0.1 tok/sec.
using:

mlx_lm.generate --model microsoft/bitnet-b1.58-2B-4T  --prompt "implement bubble sort from scratch" --max-tokens 10 --temp 0.7

with --max-tokens 100 => 81 +/-0.2 tok/sec

Q: Do you know where I can find a good example script in order to fine-tuned a Bitnet model or Falcon using mlx ?

Q: Do you know where I can find a good example script in order to fine-tuned a Bitnet model or Falcon using mlx ?

Not sure if this is implemented in MLX, but either for microsoft bitnet or Falcon-E you will need to download the pre-quantized weights (for bitnet it's in a separate repo: https://huggingface.co/microsoft/bitnet-b1.58-2B-4T-bf16 - and for Falcon-E it's on a separate revision prequantized.) Then use another implementation of BitLinear tailored for training.
The implementation of MS-Bitnet is here: https://github.com/huggingface/transformers/blob/27459025b8f77d53631f7961cc967fa659d43f7e/src/transformers/integrations/bitnet.py#L307 (you need to consider online_quant)
The implementation of Falcon-E bitnet layer is here: https://github.com/tiiuae/onebitllms/blob/main/src/onebitllms/layers/bitnet.py#L20

Ok so I posted some updates:

• General cleanup
• Removed support for Bitnet quantizing the Llama. Let's revisit that in a follow-on if we want to support that. But I want to get this landed and keep the diff self-contained
• Improved the bitnet layer. There was some inefficient casing happening there so it's even a bit faster now.

This is good to merge. Thanks for the contributions everyone!!

My pleasure!

The llama changes are there to allow for Falcon bitnet. I can make a separate PR if you don’t mind.

There was some inefficient casing happening there so it's even a bit faster now.

Thanks a lot, I noticed the speed up! (100-> 115 tok/s)

static_cast<T> fixes the bug I was having 🙌🏽

The llama changes are there to allow for Falcon bitnet. I can make a separate PR if you don’t mind.

Yes I gathered as much. You can send a PR.. I removed it from this one because the added complexity wasn't great. I don't yet have a good suggestions for how to do it in a simpler way.. but that's something that would be good to think through