Apple Silicon向けに、DFlash（論文）のネイティブMLX実装を作っています。小さなドラフトモデルがブロック拡散により16トークンを並列に生成し、ターゲットが1回の順伝播でそれらを検証します。出力はベースラインとビット単位で完全に同一です（貪欲な完全なargmax一致）。

セットアップ： M5 Max、64GB、MLX、CUDAなし。

結果

Qwen3.5-9B bf16

生成長	DFlash	ベースライン	スピードアップ
1024トークン	85 tok/s	26 tok/s	3.3x
2048トークン	80 tok/s	26 tok/s	3.1x

Qwen3.5-4B bf16

生成長	DFlash	ベースライン	スピードアップ
1024トークン	109 tok/s	41 tok/s	2.7x
2048トークン	133 tok/s	42 tok/s	3.2x

4Bは、生成が長くなるほど速くなります。文脈が伸びても、ドラフト／検証のバランスが健全に保たれるほどモデルが小さいためです。

Qwen3.5-27B 符号化（quantized）

量子化	生成長	DFlash	ベースライン	スピードアップ
8bit	1024トークン	35 tok/s	14 tok/s	2.5x
8bit	2048トークン	26 tok/s	11 tok/s	2.3x
4bit	1024トークン	44 tok/s	24 tok/s	1.9x
4bit	2048トークン	40 tok/s	23 tok/s	1.7x

8bitのほうが4bitより良いスピードアップ比になります。 int4だと検証が非常に速くなり、bf16のドラフトがボトルネックになります。int8ではドラフト／検証のバランスがより健全です。

すべての数値は生成のみです（最初のトークンから最後のトークンまで、プリフィルなし）。どのモデルでも受容率は約80〜87％です。

作ったもの

DFlashのMLX実装は存在しませんでした。ランタイムを最初から書きました。実際に数値を動かしたのは：

head_dim=256のパッチ。 Qwen3.5-9Bはhead_dim=256を使っており、MLXのsteel_attentionでは対応されていませんでした。2行のパッチで高速なSDPAパスが解放されました。

同期の省略（sync elision）。 パイプラインを、サイクルあたりGPU→CPUの同期を2回行う構成から1回へ再編しました。80+ tok/sでは、各同期のコストは約0.5msです。

QKV射影のパック。 3回のmatmul → 1回のmatmul＋分割。層ごとのカーネル起動回数を減らします。

Apple Siliconに関する学び

統一メモリではすべてが帯域（バンド幅）制約になります。これにより、推測デコードのゲームの前提が変わります：

カスタムMetalカーネル（バッチ化-GEMV、fused gated SiLU、カスタムSDPA）をいずれも試した結果、標準のMLX steel GEMMより0.5〜0.8倍遅いことが判明しました。結局、それらはすべて元に戻しました。

検証コストは4〜16トークンでほぼフラットです（57ms vs 59ms）。ボトルネックはトークン数ではなく、重みのロードです。 「信頼度が低いときは検証するトークン数を減らす」ではここでは効果がありません。

量子化モデルでは最適化の景色が反転します：ドラフト（bf16）が検証（int4/int8）より遅くなります。これはbf16の場合とは逆で、帯域制約のあるハードウェア上で量子化ターゲットに対して推測デコードを行うときの構造的な制限です。

現在取り組んでいること

ドラフト圧縮／蒸留：27B向けに、量子化ターゲットでbf16ドラフトのボトルネックを解消します。

長い文脈の安定性。 KVキャッシュの増加により、2Kトークンを超えるとスピードアップが低下します。

MoEモデル。 DFlashのドラフトはQwen3.5-35B-A3B（合計35B、アクティブは3B）で存在します。小さなモデルの検証コストと、大きなモデルの品質を両立します。

すべてまだ製作の最中です。準備ができ次第オープンソースにします。