Speculative decoding in uzu

Speculative decoding in uzu

We are releasing our speculative decoding implementation in Uzu. Initially for Qwen3.6 27B, with support for Qwen3.8 27B and Muse Glimmer coming soon. On Apple M5-series chips, we outperform MTPLX by almost 2x, and llama.cpp by over 3x at comparable quantization levels, with the strongest gains acheived on mathematical reasoning and coding tasks.

1.00

Our draft model, quantized checkpoint format, verification algorithm, and GPU kernels are co-designed from the ground up around the latest Apple M5 chips to take maximum advantage of GPU Neural Accelerators. Unlike popular speculative decoding architectures such as model-native MTP, which produce small draft chains of 3-4 tokens at a time, we use extremely agressive speculative budgets of 16-32 tokens. This enables us to use Neural Accelerator backed GEMM kernels, acheiving maximum utiliziation of hardware arithmetic throughput.

As utility of tokens in a single draft chain drops exponentially with the chain length, we utilize a tree-based speculative decoding approach. Our hybrid draft model, DFlash-Weaver, uses a small autoregressive transformer to stich top-k predictions of a parallel DFlash drafter into a tree of coherent continuations. Since Qwen3.6 is a hybrid architecture which uses Gated DeltaNet layers, we implemented highly-efficient rollback-free tree verification kernels in Metal. Dual quadratic formulation of the GDN recurrence enables us to compute the deltanet outputs over the tree of inputs in a single recurrence-free pass. We use an adaptive tree construction algorithm which generates long chains when the drafter is confident in its predictions (on domains like agentic coding), and wide branching trees when it is not (on tasks like creative writing). Instead of checkpointing the state at every node of the tree, we cache the QKV vectors and delay the state update until the next forward pass when the verified path becomes known. We use a communication-free verification procedure based on Gumbel couplings, which, despite having theoretically suboptimal acceptance rate, allows us to generate high-quality draft trees by agressive oversampling followed by heuristic-based pruning.

1.00

Unlike older generations of Apple GPUs, M5 acheives the maximum arithmetic throughput in int8 precision. Therefore, in addition to 4-bit asymmetric (Mirai-M) and 8-bit symmetric (Mirai-L) weight quantization, we use 8-bit dynamic activation quantization during inference. We apply block-diagonal random Hadamard transforms to activations before quantization, which removes the outliers and makes activation quantization virtually lossless.

Our quantization pipeline consists of two stages: We start with a post-training quantization (PTQ), for which we use a modified version of YAQA, a state-of-the-art second order quantization algorithm. We follow it with a short quantization-aware distillation stage, optimizing only the scales and zeropoints of the quantization groups. We use disjoint slices of OpenHermes-2.5 dataset for PTQ calibration and QAD training sets.

Despite being optimized for maximum decoding performance, our quantized checkpoints attain the size vs quality Pareto frontier, matching or outperforming Unsloth checkpoints at similar bits per weight rates.

Measurements

We use a fixed 1,355-token prompt to measure autoregressive output speed, input speed, and resident memory consumption. Each score is computed as a mean of three consecutive runs.

Speculative decoding output speed is computed as a macro-average of measurements collected on random samples from three datasets with default sampling parameters:
MT-Bench — conversational prompts, 80 samples
MATH-500 — mathematical reasoning, 128 samples
HumanEval — code generation, 128 samples

We evaluate the quantization quality by measuring KL divergence relative to the bfloat16 teacher on a proprietary data mixture comprised of 45% public agentic interaction logs, 30% public SFT data, and 25% private chat logs.

Models, runtime & infrastructure to
make on-device AI interactive,
ambient & continuous.

Learn more