Pull requests / #3457

#3457 [Engine] Metal: opt-in --dflash DFlash2 drafter in place of the MTP layer

open · draft · @hotocoo · 0 comments · View on GitHub

BenchmarksSetup & installServer & APIAMD / HIPNVIDIA / CUDAModels & quantsDocumentationWindows

Description

## Problem and result

On a Mac the only drafter for verify windows is the MTP layer (plus the suffix drafter). The engine cannot use a DFlash2 draft model; #3322 adds DFlash2 to the CUDA engine, and the Metal engine has nothing.

With this change, `--dflash <draft.gguf>` (opt-in, Metal `--serve`) drafts in place of `--mtp`: up to `--spec - 1` tokens a window from the outputs of target layers 5/15/25/35/45 (the mean of the four hyper-connection streams), and `--spec-min-p` stops a draft where the selector's confidence falls below it. Without `--dflash` the engine runs the same dispatches as before: every new dispatch and buffer is behind a flag only `--dflash` sets.

On an M4 Max (64 GB) with an abliterated IQ3_S target, greedy, at a 42 GB expert arena:

| | MTP (spec 3, min-p 0.7) | DFlash2 fine-tuned for the target (spec 4, min-p 0.75) |
| --- | --- | --- |
| Short prompts, decode tok/s (2 runs) | 48.1, 46.9 | 49.7, 49.7 |
| 29K prompt, cold, decode tok/s (2 runs) | 50.4, 50.0 | 51.0, 51.7 |
| 102K prompt, cold / cached follow-up, decode tok/s (1 run) | 46.8 / 44.4 | 50.5 / 45.1 |

The stock `incoai/Qwen3.8-Flash-Next-DFlash2` drafter was **slower** than MTP on this target (41.1-42.5 vs 42.1-43.7 tok/s short, arena 33). The gain comes from a drafter fine-tuned on the target's own features, not from DFlash2 alone. Greedy output is identical to MTP.

## Related discussion

- #3322 (DFlash2 on CUDA, draft) and #3304: this is the Metal counterpart. It shares no file with #3322 except `src/program/generate.cpp`, where everything is behind `#if defined(STRATA_GPU_METAL) && defined(STRATA_DFLASH)`.
- #366 (the earlier DFlash2 RFC, closed).
- No maintainer agreement on scope yet, so this stays a draft. No default changes; setup does not offer it.

## Approach

- `metal/kernels/gr.metal`: `hc_mean`, one dispatch per captured layer, writes the HC-stream mean of `R` into a feature buffer (5 x n_embd per row). `window.cpp` and `prompt_chain.inc` encode it after the hyper-connection read of layers 5/15/25/35/45 in the decode/verify window, the resident prompt chain and the streaming prompt chunks, only when `df_on_` is set.
- `src/spec/dflash_drafter.cpp`, `include/strata/spec/dflash_*.hpp`: the drafter. Two backends: llama.cpp's DFLASH graph (`libllama` from the pinned llama.cpp 3cf03257, built shared with Metal), or `STRATA_DFLASH_NATIVE=1`, Strata's own Metal kernels (`metal/kernels/dflash.metal`, `metal/src/engine/dflash.cpp`), faster per draft. Each prompt chunk and committed window injects its rows' features; a draft runs the block and the selector.
- `src/program/generate.cpp`: `--dflash` parsing and the serve loop's draft source; it refuses `--dflash` with `--mtp` and with `--spec` < 2.
- `STRATA_DFLASH_DUMP=<prefix>` (opt-in, with `--dflash`): appends each committed row's position, token and features to `<prefix>.idx` / `<prefix>.f16`. That is how the training data for the fine-tuned drafter was made; the training script is not in this PR.
- `tools/dflash_head.py`: makes the draft GGUF: a converted DFlash2 draft plus the target's `token_embd` and output head (Q8_0), optionally only the MTP draft-vocabulary rows with a `d2t` map. `metal/src/ref_quant.cpp` learns GGUF type 27 (I64) for that map.
- `metal/CMakeLists.txt`: all of it under `STRATA_DFLASH_LLAMA` (unset: nothing new is built and `--dflash` is refused).
- `docs/MAC.md`: a short `--dflash` paragraph. `bench/results/2026-10-20-metal-dflash2/`: the runs.
- The fine-tuned drafter: https://huggingface.co/hotocooo/Qwen3.8-Flash-Next-Abliterated-Uncensored-Turbo-Brilliance-50x-Reasoning-Instruct-GSQ-RCO-DFlash (Apache-2.0, from the incoai checkpoint).

## Validation

Revision: `metal-dflash2` (main 9ae93f7b + this change). The engine runs were measured on main e88ee39d + this change before two edits made for review: a per-window debug trace (`STRATA_DFLASH_TRACE`, off unless set) was removed, and the prompt path's token copy for the drafter now runs only with `--dflash`. Neither changes what runs with `--dflash`; the engine runs were not repeated on the final revision. The change applied to 9ae93f7b without conflicts. Apple M4 Max, 64 GB, macOS 26.5. llama.cpp 3cf03257 (the pinned commit), metal-cpp from setup.

**Passed:**
- Build on 9ae93f7b + this change: `cmake -S metal -B build-df -G Ninja -DCMAKE_BUILD_TYPE=Release -DMETAL_CPP_DIR=… -DSTRATA_GGML_DIR=<llama.cpp> -DSTRATA_DFLASH_LLAMA=<llama.cpp shared Metal build>` and `ninja -C build-df strata dflash_drafter_test dflash_native_parity_test`: links; warnings are existing ones plus a duplicate `-rpath` from the linker.
- The same without `-DSTRATA_DFLASH_LLAMA` (`ninja -C build-off strata`): links.
- `build-df/dflash_drafter_test t3000-q8-dv.gguf`: ok (steady draft 3.55 ms for 1 token, 9.73 ms for 7; another model server was running, so these are upper bounds).
- `build-df/dflash_native_parity_test t3000-q8-dv.gguf`: ok (native Metal drafter against an f64 host reference, 2200 context rows, 300 trials).
- Greedy, six fixed prompts, 700 output tokens (`bench/results/…/scripts/run.py`): MTP on the unpatched binary against the fine-tuned DFlash2 drafter on the patched one: 6 of 6 identical, reasoning and answer, character for character (`greedy-mtp.json`, `greedy-tuned.json`).
- Speed matrix, every config started fresh (`scripts/matrix*.py`, `matrix.tsv`, 59 rows): arenas 33/38/40/42/45/47 GB, spec 3/4/5, several `spec_min_p`, suffix drafter on/off, stock and fine-tuned drafters. 102K engine lines in `long-102k-engine-lines.txt`.

**Failed:**
- Not a code failure: arenas 45 and 47 GB page on this 64 GB Mac and decode falls to 8.8-41.1 tok/s with both drafters (kept in `matrix.tsv`).

**Not run:**
- The patched binary **without** `--dflash` against unpatched main, output by output. The MTP arms ran on an unpatched build. The no-change claim for the default path rests on reading the code: every new dispatch is behind `df_on_`, set only by `--dflash`.
- CUDA, HIP and SYCL builds. No file they compile changes except `src/program/generate.cpp`, whose additions are all inside `#if defined(STRATA_GPU_METAL) && defined(STRATA_DFLASH)`, apart from the `--dflash` option string, its usage text and its parsing.
- Sampling (temperature > 0), other Macs, the official Qwen weights or other quants, concurrent requests, the setup and `strata launch` paths.

## Limits

- Mac only, and only with `serve/server.py`; setup does not offer it.
- A drafter has to match the target's features: the stock incoai drafter does not pay on this abliterated target. The fine-tuning tools are not part of this PR.
- Few repetitions per cell: differences under about 2 tok/s on short prompts are within the spread between repeats.

Related on this site

Quick links to related pages.