Slowdown on M1 Max (32GB)

#3
by brechtm - opened

Hi,

I find that both prefill and token generation in oMLX are much slower compared to running without DFlash. I'm using the mlx-community/Qwen3.8-27B-4bit and settings as per your instructions in https://inco.ai/blog/dflash2/.

image

The chat web interface does not show tok/s when DFlash is enabled, but the built-in Throughput benchmark show this:

# Context: Code (Python)

# Single request results
Test	TTFT(ms)	TPOT(ms)	ppTPS	tgTPS	E2E(s)	Throughput	PeakMem
pp 1024 / tg 128	15996.0	426.6	64.0	2.4	70.2	16.4	16.9 GB

I pointed Pi + DeepSeek V4 Flash at the issue and it discovered that the draft acceptance rate was 0%. It also found the root cause and provided a fix:

Root cause of the 0% acceptance: the draft runs in fp16, and the DFlash2 draft forward pass produces NaN in fp16

I reproduced the bug end-to-end with the actual models on your machine. Here's the full chain:

1. Where the fp16 comes from (not the conversion β€” the runtime)

Your active DFlash profile sets dflash_draft_quant_enabled: true with weight_bits: 4 but no activation bits β†’ dflash-mlx interprets it as w4a16. In
dflash_mlx/runtime/loading.py:

  def resolve_draft_load_dtype(quant_spec, *, chip_profile=None):
      if quant_spec is None or quant_spec.act_bits == 32:
          return None
      profile = chip_profile or detect_chip()
      if profile.bf16_emulated:        # M1/M2 (arch gen 13/14)
          return mx.float16              # ← quantized draft gets cast to FP16

On your M1 Max (applegpu_g13s), the quantized draft is automatically cast to fp16 (meta logged as draft_load_dtype='float16', source
old_apple_bf16_emulation). I confirmed this by loading the draft exactly as the engine does.

2. The fp16 draft forward pass β†’ NaN

Reproduced with the real models:

  w4a16 (your config):  draft_hidden finite=False, logits all NaN, draft tokens = [248304, 248304, ...]  β†’ acceptance 0
  bf16 (no quant):      finite, draft tokens match target exactly (7/7 on a code prompt)
  w4a32 (fp32 compute): identical acceptance to bf16 [7,1,0,1,1,2] across 6 prompts

The NaN appears in fused end-to-end evaluation (breaking the graph into per-op evals stays finite) β€” a fp16 fused-kernel overflow in the DFlash2
layers. The draft is numerically unstable in fp16, so its logits are garbage and the target never accepts a single token. That's why every run logged
acceptance=0.0% and DFlash2 was ~5Γ— slower than plain (4.1 vs 19.6 tok/s).

Fix applied to your config

  • model_settings.json + model_profiles.json (dflash profile): added dflash_draft_quant_activation_bits: 32 β†’ w4a32:gs64 (quantized weights + fp32
    compute β€” verified draft_load_dtype='float32', finite, same acceptance as bf16), and pointed the draft back to the bf16 checkpoint.

Restart the oMLX server for it to take effect. Alternatives if you'd rather not restart: set dflash_draft_quant_enabled: false (bf16 compute β€” works,
but bf16 is emulated on M1 so the draft pass is slower).

With this fix and draft quantization disabled, tok/s is better but still lower than without DFlash. With draft quantization enabled and 32-bit activations, tok/s is a bit higher than without DFlash (25 tok/s).

Sign up or log in to comment