Qwen3.8-27B-DFlash2-ternary-bonsai2

A DFlash 2 draft model for prism-ml/Ternary-Bonsai-2-27B (PrismML's 2-bit, Hadamard-rotated Qwen3.8-27B). It is z-lab/Qwen3.8-27B-DFlash2 fine-tuned on the ternary model's own hidden states and its own greedy generations, so the drafter predicts what the 2-bit target will actually say rather than what the bf16 base would have said.

Same architecture and config as the z-lab checkpoint (5 Qwen3-style layers, hidden 5120, 32/8 heads, sliding window 2048, block_size 8, taps at target layers 5/19/33/47/61, rank-256 candidate selector, top-k 16, mask_token_id 248070); shares the target's embeddings and lm_head, so there are no lm_head/embedding tensors here. 81 bf16 tensors, 3.85 GB. config.json is the z-lab file, unchanged.

Runtime: NakliTechie/dflash-mlx-bonsai2 (Apple Silicon, MLX). Status: experimental; measured on one M4 Pro.

Why re-fit

The z-lab drafter was trained on bf16 residual streams with bf16 targets. On the ternary target its conditioning features drift (cosine of the tapped hidden states vs a 4-bit Qwen3.8-27B reference: 0.993 at layer 5 falling to 0.818 at layer 61, 0.834 on the concatenated fc input) and the ternary model's next-token argmax disagrees with the reference on ~23 % of tokens. Live acceptance on the ternary target was 2.5–4.4 tokens per 8-token cycle depending on the prompt, well below the ~5 the same drafter reaches on the 4-bit target.

Training

All captures and training ran on one NVIDIA L40S (AWS g6e.2xlarge, ~13 h total). Pipeline in the runtime repo under lab/aws/.

Capture. A tap-dump tool built on PrismML's llama.cpp fork runs the ternary target (PTQ1_0 GGUF, exact ternary values) on CUDA and records, per token, the residual stream at the five tap layers (int8 + fp16 row scales) and the model's own top-8 next-token ids and logprobs. The training target at every position is the ternary model's own argmax, i.e. what greedy verification accepts. The drafter's embeddings and lm_head were dequantized from the MLX pack (verified against MLX to 2–5e-4).

Round 1 (teacher-forced corpus, 2.2 M tokens). Corpus: 479 Python-stdlib files, 800 Markdown docs, 411 Gutenberg prose chunks. First fc + hidden_norm only (1 epoch, stride 32, batch 16, lr 5e-5): held-out proxy 2.85 → 2.95 tokens per cycle. Then the full 1.9 B drafter (1 epoch, stride 32, lr 2e-5, init from the adapter fit): 2.95 → 3.00.

Round 2 (not carried forward). Full fine-tune on the corpus plus 2.5 M tokens of UltraChat dialogues rendered through the pack's chat template (teacher-forced, 1 epoch). Scored on the model's own generations it landed below round 1 (3.50 vs 3.55): more human-written text moved the drafter away from the model's own stream.

Round 3 (this checkpoint). Full fine-tune from round 1 on 1,533,200 tokens of the ternary model's own greedy generations (2.5 M captured tokens; 3,836 prompts consumed of 4,879 = 3,000 UltraChat single-turn with thinking off, 800 three-message contexts, 600 with thinking on, 479 raw code prefixes; 512-token generation cap). Gen-only blocks (anchors only where all 7 drafted positions are generated tokens), 2 epochs, stride 4, window 160, lr 2e-5, 77,500 steps, 4 h 37 min. Loss stayed in a 1.8–3.0 band. Held-out proxy on generated blocks (626 blocks, 12,649 tokens): 3.55 → 3.62 tokens per cycle; per-position match 0.80 / 0.62 / 0.52 / 0.44 / 0.36 / 0.31 / 0.28.

Measured acceptance (live, MLX, Ternary-Bonsai-2-27B, M4 Pro 24 GB)

Tokens accepted per 8-token cycle (7 drafts + anchor), fixed 8-token blocks, 512 tokens, greedy, thinking on (the pack's template default), this drafter vs the shipped z-lab one, same session:

Thinking off (natural EOS): email 2.04 → 2.37, code (1024) 3.91 → 4.08, story (2048 cap) 2.26 → 2.52 — 16 % / 4 % / 12 % fewer cycles. Acceptance is deterministic across sessions (greedy); the drafter wins on every prompt measured.

End-to-end on that machine with the runtime's Metal verify kernel (~140–150 ms per cycle vs ~46 ms per plain decode step): code ~27–29 tok/s, math ~26, raw code ~33, chat ~17 against ~21.5 tok/s plain decode. Numbers and caveats in the runtime README.

Measured on NVIDIA L4 (CUDA, llama.cpp)

llama-server on PrismML's llama.cpp (prism + DFlash 2, PrismML-Eng/llama.cpp#261), target Ternary-Bonsai-2-27B-PQ2_0.gguf, this drafter's Q4_K_M GGUF, greedy, batch 1, draft length 7. Speedups are decode tok/s against the fastest plain configuration (PTQ1_0, no speculation) on the same GPU. Full tables, per-sample outputs and scripts: naklitechie/bonsai2-dflash2-bench.

• Thinking on: 1.25-1.63x (40-prompt subsets; per-prompt median 1.34-1.95x).
• Prompt-lookup speculation (ngram-mod) on the same sets: 0.92-0.95x, so the gain is drafting, not copying. HumanEval reaches 2.66x, but 64 % of each answer copies the prompt; treat it as an upper bound.
• This drafter vs the z-lab checkpoint, same machine: +2-5 % tok/s on code and math, +8 % on MT-Bench; on MT-Bench writing prompts the z-lab drafter runs at 0.84x of plain, this one at 0.98x.
• For chat, use --spec-draft-n-max 3: MT-Bench 1.49x, writing 1.26x; code and math drop about 15 %.
• Greedy output with speculation is not byte-identical to plain decode (batched verify flips near-tied tokens); accuracy differs by at most one problem per set.

Stacked with prompt lookup (--spec-type ngram-mod,draft-dflash, ngram-mod defaults): llama.cpp tries an exact 24-token n-gram match first and falls back to this drafter. When the output copies the prompt it drafts past DFlash 2's 8-token block. Same machine, same session, vs the fastest plain configuration:

Accuracy within 2 problems per set. Details and rows: stacking/ in the dataset above.

Use

git clone https://github.com/NakliTechie/dflash-mlx-bonsai2 && cd dflash-mlx-bonsai2
bash scripts/setup-bonsai2.sh          # venv, install, downloads the pack and this drafter
bash scripts/serve-bonsai2.sh          # OpenAI-compatible server on http://127.0.0.1:8790/v1

or by hand:

DFLASH_PRISM_VERIFY=v7 dflash serve \
  --model <local dir of prism-ml/Ternary-Bonsai-2-27B-mlx-2bit> \
  --draft naklitechie/Qwen3.8-27B-DFlash2-ternary-bonsai2 \
  --port 8790 --prefill-step-size 512

Use temperature: 0; speculation only engages on greedy requests. The runtime W4-quantizes the drafter at load (--draft-quant w4, the registry default); --draft-quant none keeps it bf16 with no measured acceptance difference.

With llama.cpp (CUDA). Build PrismML's llama.cpp with DFlash 2 (#261 above) and point -md at the GGUF:

llama-server -m Ternary-Bonsai-2-27B-PQ2_0.gguf -ngl 999 -fa on --jinja \
  -md Qwen3.8-27B-DFlash2-r3-Q4_K_M.gguf --spec-type ngram-mod,draft-dflash --spec-draft-n-max 7 -ngld 999

--spec-type ngram-mod,draft-dflash stacks prompt lookup in front of the drafter (best for code and editing); use --spec-type draft-dflash alone to skip it.

In the browser. LocalMind (Ternary Bonsai 2 27B, WebGPU) fetches Qwen3.8-27B-DFlash2-r3-Q4_K_M.gguf from this repo after the model loads and runs the same drafter in WGSL, weights kept Q4_K/Q6_K on the GPU (~1 GB): 1.18× on a 1,486-token code answer on an M4 Pro, output identical to plain decode. Settings → Models → Speculative decoding for Ternary Bonsai 2 27B (on by default). The port is lab/webgpu/ in the repo above.

Limits

• For the ternary target only. Trained on Ternary-Bonsai-2-27B's hidden states and generations; untested on the bf16 / 4-bit Qwen3.8-27B, where the z-lab checkpoint is the right drafter.
• Small gains, a hard ceiling. Every lever (adapter fit, corpus fine-tune, own-generation fine-tune) moved acceptance by low single digits; position-1 match sits near 0.80 whichever drafter is fitted. The ternary target's next token is less determined by its residual stream than the bf16 base's. Chat is roughly break-even end to end; code and math win.
• Thinking traces are the worst case: acceptance decays over long exploratory reasoning text. Long free-running greedy generations of the target can also fall into repetition loops (a property of the target under greedy decoding, not of the drafter).
• Needs a DFlash 2 runtime that can load the Hadamard pack. dflash-mlx-bonsai2 does. Mainline llama.cpp has DFlash 2 but cannot apply the pack's Hadamard metadata; a port of DFlash 2 onto PrismML's llama.cpp fork exists in the runtime repo's lab/leg9/patches/ but its Metal path has no small-M ternary kernel and is a net loss on Mac today.
• Qwen3.8-27B-DFlash2-r3-Q4_K_M.gguf (1.14 GB) is the same weights in llama.cpp's dflash GGUF layout (converted with lab/leg9/r3_to_gguf.py, quantized Q4_K_M), for --spec-type draft-dflash. PrismML's llama.cpp applies the target's Hadamard transforms to the borrowed embed/lm_head (PrismML #210); DFlash 2 itself is PrismML-Eng/llama.cpp#261. On CUDA it is a speed win (table above); on Metal it is not yet.

Provenance

model.safetensors = s3://skypilot-cairn-artifacts/localmind-dflash/adapter-r3/adapter_final.safetensors (full drafter, bf16, 81 tensors; the two candidate-selector codebooks carry the .weight suffix the MLX loader would add anyway). Trained 2026-09-21 on 1.53 M of the ternary model's own greedy generations (gen-only blocks, 2 epochs, stride 4, lr 2e-5) from the round-1 full fine-tune. Proxy on held-out generated blocks: 3.62 tokens/cycle (shipped z-lab drafter on the same blocks: 3.55). config.json copied from the z-lab checkpoint unchanged. Provenance: dflash-mlx-bonsai2 lab/aws/, plan LocalMind/plan/2026-09-18-dflash-bonsai2-plan.md.

Attribution

• z-lab/Qwen3.8-27B-DFlash2 and z-lab/dflash (Apache-2.0) — the method and the checkpoint this one was fine-tuned from. Paper: Chen, Liang, Liu, DFlash: Block Diffusion for Flash Speculative Decoding.
• PrismML — Ternary-Bonsai-2-27B (Apache-2.0). "Created using Bonsai by Prism ML."
• bstnxbt/dflash-mlx (Apache-2.0) — the runtime the fork builds on.
• Qwen3.8-27B (Apache-2.0) — the base model.

@misc{chen2026dflash,
  title={DFlash: Block Diffusion for Flash Speculative Decoding},
  author={Jian Chen and Yesheng Liang and Zhijian Liu},
  year={2026},
  eprint={2602.06036},
  archivePrefix={arXiv},
  primaryClass={cs.CL},
  url={https://arxiv.org/abs/2602.06036}
}