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

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 (PTQ10 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 `lmhead` 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 Q4K/Q6K 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 Q4KM), 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}
}