When Multi-Token Prediction Meets Low-Bit Large Language Models
Ning Zhang · Shihao Wang · Jinrui Zhang · Chaodong Xiao · Lei Zhang
The Hong Kong Polytechnic University
NeurIPS 2026 · Paper · Models · Quick start
BitMTP makes low-bit LLMs faster: 2×+ decoding speedups with multi-token prediction across server, consumer, and edge GPUs.
- 2026.10 🎉 We release the BitMTP code and Bonsai checkpoints.
- 2026.09 🎉 BitMTP has been accepted to NeurIPS 2026!
BitMTP couples multi-token prediction with ternary model execution. Its CUDA runtime uses width-specialized kernels and shared-KV graph replay for multi-token decoding. This repository provides Bonsai training, model conversion, and low-bit inference for H20, RTX 3090, and Jetson Orin NX.
This repository provides the Bonsai implementation with packed T2-g128 execution. Model checkpoints and their matching tokenizers are hosted on Hugging Face.
| Model | Weights | Documented platforms |
|---|---|---|
| Bonsai-1.7B SFT | SFT AR | H20, RTX 3090, Orin NX |
| Bonsai-1.7B MTP | D4 | H20, RTX 3090, Orin NX |
| Bonsai-1.7B MTP | D8 | H20, RTX 3090, Orin NX |
| Bonsai-4B SFT | SFT AR | H20 |
| Bonsai-4B MTP | D8 | H20 |
H20 and RTX 3090 evaluation profiles cover GSM8K and HumanEval+. The Jetson Orin guide provides native-T2 single-prompt generation.
Convert the QAT checkpoints with BitMTP before low-bit inference.
Clone the repository:
git clone https://github.com/PolyU-VCLab/BitMTP.git
cd BitMTPUse Python 3.10 and a CUDA development toolkit with nvcc. The following environment is for H20; for other GPUs, follow the RTX 3090 or Jetson Orin guide instead.
python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements/h20.txt
python -m pip install --no-deps -e .
python -m pip check
SM_ARCH=90 NVCC=/usr/local/cuda/bin/nvcc bash src/bonsai_sft/gpu/t2g128_kernels/compile.shSee installation for dependencies and RTX 3090 for its separate environment and BF16 draft path.
Download Bonsai-1.7B D8, convert it, and run hybrid decoding on H20:
python -m bonsai_sft.download_model \
--model bonsai_1p7_d8 --output weights/BitMTP-Bonsai-1.7B-D8
python -m bonsai_sft.gpu.convert_bonsai_mtp_to_gpu \
--src weights/BitMTP-Bonsai-1.7B-D8 \
--out weights/BitMTP-Bonsai-1.7B-D8-packed --quant ternary
python -m bonsai_sft.gpu.run_generation \
--ckpt weights/BitMTP-Bonsai-1.7B-D8-packed \
--hf-ckpt weights/BitMTP-Bonsai-1.7B-D8 \
--block-size 8 --mode hybrid --threshold 0.9 --min-chain 6 \
--draft-precision fp8 --verify-precision native \
--chat --prompt 'Write a short greeting.' --max-new-tokens 128For Bonsai-1.7B D4, use D4 weights, --block-size 4, and --min-chain 3. See generation for all decoding modes.
| Task | Guide |
|---|---|
| Train SFT/MTP models on your JSONL data | Training |
| Evaluate on GSM8K and HumanEval+ | Evaluation |
| Run software regression tests | Tests |
Training examples are in configs/training/, and evaluation profiles are in configs/evaluation/.
src/bonsai_sft/ Training, model export, and CUDA decoding
configs/ Training and evaluation settings
benchmarks/ GSM8K and HumanEval+ evaluation
scripts/ Training launcher
tests/ Software regression tests
docs/ Installation, usage guides, and method overview
requirements/ Platform and CPU-scoring dependencies
docker/ Isolated HumanEval+ CPU scorer
licenses/ Third-party license texts
BitMTP code is licensed under Apache-2.0. Third-party components retain their respective copyrights and licenses; see THIRD_PARTY_NOTICES.md. Model weights and datasets retain their respective terms.
We thank the Bonsai, SDLM, and BitNet / bitnet.cpp teams and contributors, and the broader open-source community advancing low-bit language models and efficient inference.