[KongchangAI]
· 2 min read· 1,247 words

BitMTP: Co-Accelerating Multi-Token Prediction and Low-Bit LLMs

BitMTP: Co-Accelerating Multi-Token Prediction and Low-Bit LLMs

BitMTP uses joint algorithm-system design to make low-bit models and multi-token prediction truly complementary, achieving over 2x decoding speedup.

BitMTP is an inference acceleration solution from PolyU-VCLab that addresses the "1+1<2" problem when combining ultra-low-bit quantization with Multi-Token Prediction. It trains a single model with joint next-token and MTP objectives, eliminating the need for a separate draft model, and introduces accepted-prefix alignment to keep KV boundaries and CUDA Graph replay states consistent. It achieves 2.0-2.4x average speedup on H20, RTX 3090, and Jetson Orin, with up to 3.3x on MBPP+.

BitMTP paper title

With LLM deployment costs remaining high, two technical routes are racing forward on separate tracks: one toward lighter models—ultra-low-bit backbone—and the other toward faster generation—Multi-Token Prediction (MTP). The former cuts weight storage, memory traffic, and linear-layer compute overhead, while the latter compresses serial decoding steps by predicting multiple future tokens in parallel. The two appear naturally complementary, but simply stitching them together fails to deliver the promised speedup at equivalent quality. BitMTP, proposed by the research team at PolyU-VCLab, takes direct aim at this awkward "1+1<2" situation.

Why Simple Stacking Doesn't Work

To understand BitMTP's value, one must first see the two critical flaws in directly combining MTP with a low-bit backbone.

The first lies in the drafting stage. MTP's typical approach has the model draft multiple subsequent tokens at once from the current position, but this kind of "masked drafting" lacks the intermediate token inputs that naturally exist during next-token training. In other words, the draft model guesses the future without seeing step-by-step context, so acceptance rates naturally suffer. The second lies at the system level. Even if the number of prediction steps decreases, wider computation and causal key-value (KV) updates erode the gains from step reduction—fewer steps, but each step is more expensive, and the math doesn't work out.

Together, these two points explain why the direct combination of "low-bit + MTP" struggles to match either method alone in decoding quality. BitMTP's entry point is algorithm-system co-design, rather than patching a single layer.

Multi-Token Prediction (MTP) is a training objective and inference acceleration technique: the model predicts the next token while also predicting multiple subsequent tokens in parallel. During training, this amounts to adding auxiliary prediction heads, letting the model learn longer-range future dependencies; during inference, these predictions can be proposed as a "draft" in one shot, then verified and multiple tokens accepted in parallel, reducing serial decoding steps. MTP's speedup depends on the draft's acceptance rate, which in turn depends heavily on whether the context visible during drafting is sufficient.

Ultra-low-bit quantization refers to compressing model weights to 1-2 bits (e.g., BitNet b1.58's ternary weights) to drastically reduce storage, memory bandwidth, and matrix multiply-add overhead. But the limited numerical precision of low-bit backbones imposes extra requirements on draft quality, KV cache representation, and kernel implementation, and combining with MTP requires dedicated design to preserve quality.

One Model, Two Faces: AR Generation and Block-Causal Drafting

BitMTP's core mechanism is a joint next-token prediction and MTP objective. This training objective equips a single model with two capabilities simultaneously: autoregressive (AR) generation and block-causal drafting. The direct benefit is bidding farewell to the "draft model + verification model" dual-model architecture—no more maintaining a separate draft model and its alignment with the main model, reducing both memory footprint and deployment complexity.

For decoding strategy, BitMTP offers three selectable modes: confidence-only decoding, speculative-only decoding, and hybrid decoding. The three correspond to different quality-throughput trade-off points, and users can flexibly switch based on how quality-sensitive the task is. This "one model, multiple operating points" design philosophy effectively downgrades decoding strategy from an architectural decision to a runtime configuration—quite friendly for real-world deployment.

The System-Level Contract: Prefix Alignment and Graph Reuse

Beyond the algorithm, BitMTP introduces a key concept in system implementation—accepted-prefix alignment. It requires three states to remain strictly consistent when continuing generation: the committed prefix, the materialized causal KV boundary, and the CUDA Graph replay state.

This "contract" matters because speculative decoding inherently involves "partial acceptance" intermediate states: after verification, only the first k tokens of a draft may be accepted, and if the KV cache and graph capture state are not updated synchronously, continuation will produce errors. Using this contract as a baseline, BitMTP further implements width-specialized low-bit execution and reuses captured computation graphs over shared KV storage. Width specialization means selecting appropriate low-bit kernels for different batch or sequence widths, avoiding the efficiency loss of a one-size-fits-all approach; graph reuse directly eliminates repeated graph capture overhead.

CUDA Graph is a technology provided by NVIDIA that captures multiple GPU kernel launches into a single static execution graph, significantly reducing kernel launch overhead and CPU scheduling overhead. In speculative decoding, the computation graph shape (batch, sequence width) may change at each decoding step, so graphs need to be captured separately for different widths, and graph state must strictly correspond to the actual KV cache and accepted prefix during reuse. BitMTP's "accepted-prefix alignment" is precisely designed to prevent misalignment among graph replay state, KV boundary, and committed prefix in partial-acceptance scenarios, enabling safe reuse of captured computation graphs.

Width-specialized low-bit execution refers to selecting specially tuned low-bit matrix multiplication kernels for different batch sizes or sequence widths, rather than using the same set of kernels for all shapes. Low-bit kernel performance is shape-sensitive, and width specialization avoids inefficient memory access and idle compute units caused by a one-size-fits-all approach.

Measured Data: Speedup Across Hardware and Models

According to evaluation results disclosed in the paper, on NVIDIA H20 with hybrid decoding, BitMTP achieves a five-task average speedup of 2.2x on BitNet b1.58-2B, and 2.1x to 2.4x on Bonsai series models ranging from 1.7B to 4B, with comparable task quality. On RTX 3090 and Jetson Orin, five-task average throughput improvement reaches up to 2.0x, with the MBPP+ single task reaching 3.3x.

The coverage of these numbers is noteworthy: they include both data-center GPUs and consumer-grade graphics cards and edge devices; they validate both BitNet, a representative low-bit backbone, and different scales of the Bonsai series. Generalization across backbones, model sizes, and GPU types is an important dimension for judging whether an acceleration technique can be deployed in practice. The research team has open-sourced the code in the PolyU-VCLab/BitMTP repository on GitHub, facilitating reproduction and further development.

A Few Observations

BitMTP's insight goes beyond a set of speedup ratios. It shows that there are structural coupling points between the two efficiency routes of low-bit quantization and MTP—draft quality, KV boundary, and graph state are interlocked, and only joint design can unlock the combined benefits. For teams focused on edge deployment and inference cost optimization, this "algorithm-system co-design" approach is more valuable as a reference than any single trick. Of course, the speedup ratios reported in the paper were all measured under specific model and task combinations; real-world business gains still need to be evaluated against one's own workloads, and the tuning space for strategies like hybrid decoding is also worth further exploration.

One-line Summary

BitMTP uses joint algorithm-system design to make low-bit models and multi-token prediction truly complementary, achieving over 2x decoding speedup.

Paragraph Summary

BitMTP is an inference acceleration solution proposed by PolyU-VCLab, aimed at solving the "1+1<2" problem when ultra-low-bit quantization and Multi-Token Prediction (MTP) are directly combined. The article identifies two critical flaws: masked drafting lacks intermediate token inputs, lowering acceptance rates; and wider computation and causal KV updates offset the gains from step reduction. BitMTP trains a single model with a joint next-token prediction and MTP objective, giving it both autoregressive generation and block-causal drafting capabilities without a separate draft model; for decoding, it offers confidence-only, speculative, and hybrid modes. At the system level, it introduces an accepted-prefix alignment contract to ensure consistency among the committed prefix, materialized causal KV boundary, and CUDA Graph replay state, and implements width-specialized low-bit execution and graph reuse. Measured on H20, RTX 3090, and Jetson Orin, it achieves 2.0-2.4x average speedup, with the MBPP+ single task reaching 3.3x; the source code is open-sourced.

#BitMTP#Multi-Token Prediction#low-bit quantization#LLM inference acceleration#speculative decoding#CUDA Graph#BitNet
Share:

Related articles