Feature request: BitNet 1.58-bit ternary inference on ROCm (gfx1151)

Request

Enable BitNet 1.58-bit (ternary weight) inference on ROCm builds. The architecture is listed as supported in the README but the pre-built gfx1151 binary fails to load BitNet models.

What We Tested

Binary: b1004-tech-preview (gfx1151)
Model: mlx-community/bitnet-b1.58-2B-4T
Result: warmup failed — model loads but inference crashes

Also tested mlx-community/Falcon-E-3B-Instruct-1.58bit — same result.

Why This Matters

BitNet 1.58-bit uses ternary weights {-1, 0, 1}. On unified memory hardware like Strix Halo (128GB), 1-bit models would:

• Dramatically reduce memory bandwidth pressure — the main bottleneck on unified memory
• Enable much larger effective model sizes — a 70B 1-bit model fits where a 70B FP16 cannot
• Pair with MLX speed advantage — MLX is already 29-85% faster than vLLM on this hardware. Adding 1-bit would compound that.

Microsoft open-sourced BitNet. PrismML has Bonsai 1-bit MLX models (Apple Silicon). The architecture support is in the codebase — the ternary matmul kernels just need the ROCm/HIP path.

Available Models

mlx-community/bitnet-b1.58-2B-4T           — Microsoft official
mlx-community/Falcon-E-3B-Instruct-1.58bit — Falcon extreme quant
prism-ml/Bonsai-8B-mlx-1bit                — PrismML (Apple Silicon)
prism-ml/Bonsai-4B-mlx-1bit
prism-ml/Bonsai-1.7B-mlx-1bit

Test Environment

Hardware: AMD Strix Halo, 128GB unified, gfx1151
Binary:   b1004-tech-preview
OS:       CachyOS (Arch), kernel 7.0.0-1-mainline

Context

We have comprehensive MLX benchmarks on this hardware — 6 models passing at 21-151 tok/s (4-bit). Full results: https://github.com/stampby/bleeding-edge

Happy to test 1-bit builds when available.

Update: Full BitNet Implementation + Build Findings

We've written a complete BitNet 1.58-bit implementation and documented everything we found. Branch: stampby/lemon-mlx-engine:bitnet-support

Implementation (302 lines, 4 files)

Files:

• include/mlx-lm/llm/models/bitnet.h — Full model header: BitNetAttention, BitNetMLP, BitNetTransformerBlock, BitNetModelInner, BitNetModel
• src/llm/models/bitnet.cpp — Implementation with uint8 ternary dequantization, relu_squared activation, sub-layer norms
• src/llm/llm_factory.cpp — Factory registration in both loader and type registries
• CMakeLists.txt — Source file addition

Architecture:

• Llama variant with three key differences:
  1. relu_squared activation instead of silu
  2. Sub-layer norms (attn_sub_norm before o_proj, ffn_sub_norm before down_proj)
  3. Ternary weights packed as uint8 (4 values per byte), dequantized at sanitize time
• Config reuses LlamaConfiguration since fields are identical
• Supports both uint8 packed weights (mlx-community) and pre-dequantized bf16 (microsoft)

Weight mapping (bitnet-b1.58-2B-4T):

hidden_size=2560, intermediate_size=6912, num_heads=20, num_kv_heads=5, head_dim=128
q_proj:    [640, 2560]  uint8 → [2560, 2560]  float16
k_proj:    [160, 2560]  uint8 → [640, 2560]   float16
v_proj:    [160, 2560]  uint8 → [640, 2560]   float16
o_proj:    [640, 2560]  uint8 → [2560, 2560]  float16
gate_proj: [1728, 2560] uint8 → [6912, 2560]  float16
up_proj:   [1728, 2560] uint8 → [6912, 2560]  float16
down_proj: [640, 6912]  uint8 → [2560, 6912]  float16

Current Status

• Model loads ✅ — sanitize/dequantize/weight_map all work
• Diagnose passes ✅ — GPU ops, matmul, RMSNorm, RoPE all green
• Forward pass ❌ — Reshape error during inference ("size 2560 into garbage shape")

The reshape error occurs during the attention forward pass. We've verified all weight shapes are correct post-dequantization, but something in the compute graph produces an unexpected tensor shape at runtime.

Build Environment Findings

Your CI builds on Ubuntu 22.04 with ROCm clang from /opt/rocm/lib/llvm/bin/clang++. We tested multiple build paths:

Environment	Result
Arch Linux + GCC 15.2.1	❌ std::optional constraint satisfaction bug in GCC 15
Arch Linux + ROCm clang + Arch libstdc++	❌ Same bug — clang uses GCC 15 headers
Ubuntu 22.04 container + system ROCm (Arch-built)	❌ GLIBC mismatch (ROCm needs 2.38, Ubuntu has 2.35)
Ubuntu 24.04 container + system ROCm	❌ Missing libxml2.so.16 (Arch vs Ubuntu soname)
Ubuntu 22.04 container + ROCm 7.2.1 from apt	❌ HIP indexing.hip structured binding error
Pre-built b1004-tech-preview binary	❌ "Unsupported model type: bitnet"

Bottom line: We can't build locally on Arch/CachyOS with current GCC 15. Your CI (Ubuntu 22.04 + ROCm apt + clang) is the right path. The HIP indexing.hip error we hit may be a dep version mismatch on our end.

What We Need

1. Try building our branch in your CI — the BitNet code itself compiles fine (tested with GCC 11/13). The issue is ROCm backend deps.
2. Help debugging the forward pass reshape error — we suspect it's in how linear_forward interacts with the dequantized weights, possibly the QuantizedWeightRegistry registering BitNet's .weight_scale keys incorrectly.
3. Or: if you prefer, just point us at the right approach and we'll adapt our implementation.

Test Hardware

AMD Strix Halo, 128GB unified, Radeon 8060S (gfx1151)
CachyOS (Arch), kernel 7.0.0-1-mainline
System ROCm 7.2.1

Context

We have all three inference backends running simultaneously on this hardware right now:

• MLX Engine (ROCm) — 21-151 tok/s on 4-bit models
• Prism llama.cpp (Vulkan) — 72 tok/s on TQ1_0 1-bit
• lemond/FLM (NPU) — 11 tok/s on Qwen3-8B

BitNet on MLX would complete the picture — 1-bit inference on ROCm with HuggingFace-native models, no GGUF conversion needed.

Happy to test any builds or patches.

[Geramy]

Update: Full BitNet Implementation + Build Findings

We've written a complete BitNet 1.58-bit implementation and documented everything we found. Branch: stampby/lemon-mlx-engine:bitnet-support

Implementation (302 lines, 4 files)

Files:

• include/mlx-lm/llm/models/bitnet.h — Full model header: BitNetAttention, BitNetMLP, BitNetTransformerBlock, BitNetModelInner, BitNetModel
• src/llm/models/bitnet.cpp — Implementation with uint8 ternary dequantization, relu_squared activation, sub-layer norms
• src/llm/llm_factory.cpp — Factory registration in both loader and type registries
• CMakeLists.txt — Source file addition

Architecture:

• Llama variant with three key differences:

  1. relu_squared activation instead of silu
  2. Sub-layer norms (attn_sub_norm before o_proj, ffn_sub_norm before down_proj)
  3. Ternary weights packed as uint8 (4 values per byte), dequantized at sanitize time

• Config reuses LlamaConfiguration since fields are identical

• Supports both uint8 packed weights (mlx-community) and pre-dequantized bf16 (microsoft)

Weight mapping (bitnet-b1.58-2B-4T):

hidden_size=2560, intermediate_size=6912, num_heads=20, num_kv_heads=5, head_dim=128
q_proj:    [640, 2560]  uint8 → [2560, 2560]  float16
k_proj:    [160, 2560]  uint8 → [640, 2560]   float16
v_proj:    [160, 2560]  uint8 → [640, 2560]   float16
o_proj:    [640, 2560]  uint8 → [2560, 2560]  float16
gate_proj: [1728, 2560] uint8 → [6912, 2560]  float16
up_proj:   [1728, 2560] uint8 → [6912, 2560]  float16
down_proj: [640, 6912]  uint8 → [2560, 6912]  float16

Current Status

• Model loads ✅ — sanitize/dequantize/weight_map all work
• Diagnose passes ✅ — GPU ops, matmul, RMSNorm, RoPE all green
• Forward pass ❌ — Reshape error during inference ("size 2560 into garbage shape")

The reshape error occurs during the attention forward pass. We've verified all weight shapes are correct post-dequantization, but something in the compute graph produces an unexpected tensor shape at runtime.

Build Environment Findings

Your CI builds on Ubuntu 22.04 with ROCm clang from /opt/rocm/lib/llvm/bin/clang++. We tested multiple build paths:

Environment Result
Arch Linux + GCC 15.2.1 ❌ std::optional constraint satisfaction bug in GCC 15
Arch Linux + ROCm clang + Arch libstdc++ ❌ Same bug — clang uses GCC 15 headers
Ubuntu 22.04 container + system ROCm (Arch-built) ❌ GLIBC mismatch (ROCm needs 2.38, Ubuntu has 2.35)
Ubuntu 24.04 container + system ROCm ❌ Missing libxml2.so.16 (Arch vs Ubuntu soname)
Ubuntu 22.04 container + ROCm 7.2.1 from apt ❌ HIP indexing.hip structured binding error
Pre-built b1004-tech-preview binary ❌ "Unsupported model type: bitnet"
Bottom line: We can't build locally on Arch/CachyOS with current GCC 15. Your CI (Ubuntu 22.04 + ROCm apt + clang) is the right path. The HIP indexing.hip error we hit may be a dep version mismatch on our end.

What We Need

1. Try building our branch in your CI — the BitNet code itself compiles fine (tested with GCC 11/13). The issue is ROCm backend deps.
2. Help debugging the forward pass reshape error — we suspect it's in how linear_forward interacts with the dequantized weights, possibly the QuantizedWeightRegistry registering BitNet's .weight_scale keys incorrectly.
3. Or: if you prefer, just point us at the right approach and we'll adapt our implementation.

Test Hardware

AMD Strix Halo, 128GB unified, Radeon 8060S (gfx1151)
CachyOS (Arch), kernel 7.0.0-1-mainline
System ROCm 7.2.1

Context

We have all three inference backends running simultaneously on this hardware right now:

• MLX Engine (ROCm) — 21-151 tok/s on 4-bit models
• Prism llama.cpp (Vulkan) — 72 tok/s on TQ1_0 1-bit
• lemond/FLM (NPU) — 11 tok/s on Qwen3-8B

BitNet on MLX would complete the picture — 1-bit inference on ROCm with HuggingFace-native models, no GGUF conversion needed.

Happy to test any builds or patches.

There is a thing about building MLX, you have to use the correct compiler otherwise ont he fly kernel compilation names dont match. I believe you have to use the same compiler as hipcc uses. I could of sworn it was GCC-15 otherwise it wont be able to run at-time compiled kernels.

Good catch — that explains why we couldn't get a clean build on Arch. GCC 15.2.1 has the std::optional constraint bug, and since hipcc needs GCC-15 for JIT kernel name matching, we're stuck until GCC 15.3 or a patched 15.2.

The pre-built b1004-tech-preview binary works perfectly for 4-bit models (151 tok/s on Qwen3-0.6B), but it doesn't have the bitnet model type registered.

Would you be able to kick our branch through your CI? The BitNet code itself compiles clean on GCC 11/13 — it's only the ROCm backend deps that need your Ubuntu 22.04 + ROCm clang setup.

Branch: stampby/lemon-mlx-engine:bitnet-support

If it builds, we can test the binary on our Strix Halo immediately. The forward pass reshape error might even be a compiler-specific issue that goes away with the right toolchain.

[Geramy]

Good catch — that explains why we couldn't get a clean build on Arch. GCC 15.2.1 has the std::optional constraint bug, and since hipcc needs GCC-15 for JIT kernel name matching, we're stuck until GCC 15.3 or a patched 15.2.

The pre-built b1004-tech-preview binary works perfectly for 4-bit models (151 tok/s on Qwen3-0.6B), but it doesn't have the bitnet model type registered.

Would you be able to kick our branch through your CI? The BitNet code itself compiles clean on GCC 11/13 — it's only the ROCm backend deps that need your Ubuntu 22.04 + ROCm clang setup.

Branch: stampby/lemon-mlx-engine:bitnet-support

If it builds, we can test the binary on our Strix Halo immediately. The forward pass reshape error might even be a compiler-specific issue that goes away with the right toolchain.

You should be able to submit a PR into my repo.

[bong-water-water-bong]

PR #41 adds native BitNet 1.58-bit support to lemon-mlx-engine. This is a reworked version of the closed PR #12, ported to the current codebase with two bugs fixed:

1. Dangling reference — the same bug as PR Fix SIGFPE crash from dangling config reference in all model constructors #40. BitNetAttention::args_ stored a reference to the constructor parameter, which dangled after load_typed_model() returned.

2. Dequantization ordering — the original stack + reshape interleaved weight rows incorrectly. Fixed to use concatenate along axis 0, matching the HuggingFace transformers reference.

Tested on Strix Halo (gfx1151)

$ echo 'The capital of France is' | ./chat ~/models/bitnet-2b --raw
Assistant: Paris, and it is known for its rich history, famous landmarks...

Model loads at 5.9 GB active memory, generates coherent text.

@stampby — this is your original PR #12 work, reworked for the current architecture. The "size 2560 into garbage shape" reshape error you hit was caused by the dangling reference (freed stack memory producing garbage num_attention_heads), and the forward pass would have produced garbage output from the dequantization ordering bug even if it hadn't crashed.

[bong-water-water-bong]

Clarification on scope: PR #41 covers 1 of the 5 models listed in this issue.

Model	model_type	Status
mlx-community/bitnet-b1.58-2B-4T	bitnet	✅ PR #41
mlx-community/Falcon-E-3B-Instruct-1.58bit	llama	❌ Not covered — BitNet quant on standard Llama arch
prism-ml/Bonsai-8B-mlx-1bit	qwen3	❌ Not covered — different format (affine quant)
prism-ml/Bonsai-4B-mlx-1bit	qwen3	❌ Not covered
prism-ml/Bonsai-1.7B-mlx-1bit	qwen3	❌ Not covered

The Bonsai models use a completely different quantization format — affine quantization with per-group .scales and .biases (not uint8 ternary packing). That's tracked separately by issue #11 ("Unsupported bits for affine_dequantize").

The Falcon-E model uses BitNet-style ternary quantization (quant_method: "bitnet") but on a standard Llama architecture (model_type: "llama", hidden_act: "silu", no sub-layer norms). Supporting it would require adding the ternary dequantization path to the Llama loader itself, rather than a separate BitNet model.

[bong-water-water-bong]

All gaps now closed — PR #41

PR #41 now supports all 5 models listed in this issue:

Model	model_type	Quant format	Status
bitnet-b1.58-2B-4T	bitnet (relu2)	uint8 ternary + weight_scale	✅ Working
Falcon-E-3B-Instruct-1.58bit	bitnet (silu)	uint8 ternary + weight_scale	⚠️ Loads and runs; model checkpoint is broken (confirmed: HF transformers also produces garbage)
Bonsai-1.7B-mlx-1bit	qwen3	1-bit affine	✅ Working
Bonsai-4B-mlx-1bit	qwen3	1-bit affine	✅ Working
Bonsai-8B-mlx-1bit	qwen3	1-bit affine	✅ Working (same code path)

What was implemented

1. BitNet ternary dequantization — shared dequantize_bitnet_weight in bitnet_utils.h, used by both BitNetModel and LlamaModel's sanitize_impl
2. Adaptive BitNetModel — uses relu² + sub_norms when hidden_act=relu2, falls back to silu for Falcon-E
3. 1-bit affine dequantization (also closes Unsupported bits for affine_dequantize - prism-ml/Bonsai-8B-mlx-1b #11) — dequantize_1bit() in quantize_utils.cpp since MLX GPU affine_dequantize kernel doesn't support 1-bit
4. YaRN rope scaling — Qwen3Attention now handles rope_type=yarn (previously only linear)

Verified on both GPU architectures

• gfx1151 (Strix Halo): BitNet ✅, Bonsai 1.7B ✅, Bonsai 4B ✅
• gfx1201 (RX 9070 XT): BitNet ✅ (OOM on long gen), Bonsai 1.7B ✅ (104 tok/s)

About Falcon-E

The mlx-community/Falcon-E-3B-Instruct-1.58bit quantized checkpoint is broken. Verified by loading it in HuggingFace transformers (CPU), which also produces garbage. The original unquantized tiiuae/Falcon-E-3B-Instruct works correctly. Our code handles it properly — the issue is with the quantized weights themselves, not our implementation.