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

[stampby]

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.

[stampby]

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.

[stampby]

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.