M1 Max: Long-context API requests are slower with DFlash than basic decoding over Qwen3.6-35B-A3B

[hdmi]

I am seeing a large gap between the local smoke benchmark and real agent API calls using the dame target/draft pair

Summary

On real chat/agent requests with ~8.4k prompt tokens, dflash serve is much slower than baseline autoregresive (AR) mode during decode.
The long prefill time is expected (because of the bigger prompt), but post prefill decode is also much worse than plain decoding.
In the mac m1 max, plain AR is roughly 50-60 tok/s and DFlash stays around 15-20 tok/s.

Environment

• dflash-mlx 0.1.5.1
• mlx 0.31.2
• mlx-lm 0.31.3
• Python 3.14.3
• Apple M1 Max, 64 GB
• target: mlx-community/Qwen3.6-35B-A3B-4bit
• draft: z-lab/Qwen3.6-35B-A3B-DFlash

What I ran

Smoke benchmark looked good:

• ~102 prompt tokens
• dflash tok/s: 110.97
• acceptance: 0.91

But real server/API diagnostics are very different.

1. verify_mode=auto

From .artifacts/dflash/diagnostics/20260511-165200-serve-basic

• request 1: 8389 prompt tokens, 246 generated, decode_tok_s=15.15, acceptance=0.561, tokens_per_cycle=2.28
• request 2: 8392 prompt tokens, 332 generated, decode_tok_s=19.60, acceptance=0.663, tokens_per_cycle=2.96

2. verify_mode=adaptive

From .artifacts/dflash/diagnostics/20260511-172625-serve-basic

• request 1: 8392 prompt tokens, 332 generated, decode_tok_s=20.14, acceptance=0.663, tokens_per_cycle=2.96
• request 2: 8389 prompt tokens, 206 generated, decode_tok_s=17.75, acceptance=0.597, tokens_per_cycle=2.48
• adaptive did engage on one request: adaptive_block_reductions=1, adaptive_block_cycles=12, adaptive_block_min=4

better, but still far below plain AR

3. verify_mode=adaptive + non-quantized draft

From .artifacts/dflash/diagnostics/20260511-172930-serve-basic

Serve command included --draft-quant none.

Results:

• request 1: 8389 prompt tokens, 432 generated, decode_tok_s=20.07, acceptance=0.669, tokens_per_cycle=3.02
• request 2: 8392 prompt tokens, 324 generated, decode_tok_s=19.33, acceptance=0.660, tokens_per_cycle=2.95

Slightly better drafting, but not enough and memory increased.

What I checked already

I do not see an obvious serve-vs-benchmark config mismatch in the core runtime:

• profile: balanced
• split-sdpa enabled
• draft_sink_size=64
• draft_window_size=1024
• verify_len_cap=0
• target_fa_window=0
• draft block_size=16

Question

Is this expected for long-context real chat/agent workloads?
If not, what should I try next? are there any recommended serve settings for long context agentic traffic?

[bstnxbt]

Hey, thanks for the detailed report, numbers, three variants, clean config, that's exactly what makes this actionable.

Both target and draft default to W4 on Qwen3.6-35B-A3B (registry sets draft_quant="w4", target is already 4-bit in the HF name). So your runs 1-2 are W4/W4, run 3 is W4 target / bf16 draft.

My read: this is likely architectural, not a config bug. On MoE-A3B at 4-bit, only ~3B params are active per token, so baseline AR is already cheap (memory-bandwidth bound on ~3B). Speculative decoding wins when target cost/token >> draft cost/token; here that gap is small and per-cycle overhead (verify batched on 16 positions over 8.4k KV, MoE expert routing per-token) eats what's left. Acceptance dropping 0.91 → ~0.6 on real prompts confirms the smoke bench isn't a valid baseline for this workload.

Before I touch anything in the repo, two tests would isolate the cause:

1. Run mlx_lm.generate on the exact same 8.4k prompt and post the decode tok/s. I want your baseline on your machine, not a generic figure.
2. Rerun dflash serve on the same 8.4k prompt with a dense target instead of MoE-A3B —mlx-community/Qwen3.6-27B-4bit is a good one. Default draft will be picked up automatically.

If dflash wins on the dense target and loses on the MoE one, the issue is "MoE-A3B + long context" and the fi is a doc/default change. If it loses on both, there's a real bug and I'll dig.

[hdmi]

Quick update with exact-body AR controls and the dense-target check

Method

• Sent same real agent request to all tests
• DFlash runs used tools.benchmarks.agentic_trace replay on that captured body against a fresh server
• For plain AR baseline i have used a local helper that uses mlx_lm.server tokenization, then runs plain mlx_lm.generate
• For apples-to-apples decode throughput, I used fixed-length --no-eos AR runs matching the DFlash completion length

Qwen 3.6

Target	Path	Prompt	Gen	Prefill tok/s	Decode tok/s	TTFT / prefill	Wall	Accept	TPC
35B-A3B-4bit	AR (mlx_lm.generate, fixed 324)	8392	324	463.97	56.74	18.09s	23.98s	-	-
35B-A3B-4bit + DFlash	DFlash replay (adaptive, bf16 draft)	8392	324	413.70	19.34	20.28s	37.04s	0.660	2.95
35B-A3B-4bit + DFlash	DFlash replay, cache repeat	8392	323-324	-	19.7-21.2	0.24-0.75s	15.5-17.2s	-	-
27B-4bit	AR (mlx_lm.generate, fixed 430)	8392	430	83.10	16.69	100.99s	126.94s	-	-
27B-4bit + DFlash	DFlash replay (adaptive, W4/W4)	8390	430	79.25	5.97	105.87s	177.87s	0.677	3.09

Notes:

• On the MoE target, prefix cache helps TTFT on exact repeats, but decode stays ~20 tok/s
• On the dense target, DFlash also loses badly vs AR on the same captured 8.4k body
• I also checked dense natural-stop AR: untouched prompt stopped after 2 tokens; with a tiny “do not use any tool, answer directly” suffix it answered directly at ~15.5-15.7 tok/s, consistent with the fixed-length AR control

So this does not look specific to MoE-A3B + long context. DFlash is also materially slower on the dense 27B target on the same long prompt

Could there be something pathological in the DFlash verify / replay / commit loop under long context, even with acceptance around 0.68?

[bstnxbt]

Thanks for the report. After looking at this and
several similar reports across the MLX ecosystem (omlx
#770, #604, #880; mlx-engine #103, #113; mlx-examples
#1281), I think there are three compounding causes for
DFlash underperforming on M1/M2 (and to a smaller
degree M3/M4), plus one Ultra-specific hypothesis.
Posting it here as a working theory + concrete fix
tiers so we can collaborate on what's worth doing.

Causes

1. BF16 is software-emulated on Apple7 / Apple8 (M1,
   M2). MLX's _MLX_BFloat16 struct unpacks every bf16
   value to fp32, computes, and repacks. Tax is around
   15-30% on every kernel that runs bf16 on these GPUs.
   DFlash drafts are bf16 by default, so we eat this on
   the hottest path.

2. No NAX matrix path before M5. M1/M2/M3/M4 fall
   back on simdgroup_matrix 8×8 tiles, 2-4× smaller than
   the M5 NAX path. Less throughput per dispatch.

3. MLX 2-pass SDPA ladder mis-tuned for g-class
   GPUs (base + Pro). The block heuristic caps at 32/64
   blocks and doesn't scale with N, so long-context
   attention is left on the table. DFlash ports it
   verbatim, so we inherit the regression.

4. (Hypothesis, Ultra-specific) UltraFusion
   small-buffer latency on cross-die dispatches.
   Single-sourced for now (ANE bandwidth report + #492),
   needs a cross-die microbench before I'd believe it.

Fix tiers, ranked by effort

Tier 1 — force FP16 draft on Apple7 / Apple8 (1-2
days). Detect arch_gen ∈ {13, 14} at load time,
cast the draft model to fp16 instead of bf16. Per omlx
#880 ("FP16 strikes again on M1/M2"), this is the
single biggest unblock for M1/M2 hardware.

Tier 2 — retune _compute_sdpa_2pass_blocks for
g-class + propagate the fp16 cast into verify_qmm for
Apple7/8 (3-5 days).

Where I'm stuck

I'm on an M5 Max, so I can't measure any of this on the
affected hardware. I can draft the Tier 1 patch (arch
detection + fp16 draft cast at load time, one place,
not scattered) blind, but the speedup gate has to be
measured on your machine.

Two ways to proceed:

• Option A: I draft the Tier 1 patch on a branch,
  you run dflash benchmark --model <your-target> --draft <your-draft> and dflash serve on your actual prompt
  before/after, post numbers here. If Tier 1 doesn't
  move the floor, Tier 2/3 wait.
• Option B: You try the FP16 draft cast yourself —
  happy to walk through the load path
  (dflash_mlx/runtime/loading.py and
  dflash_mlx/draft_backend.py are the two spots) and
  what to flip.

[hdmi]

Implemented a local Tier 1 patch on my M1 Max and reran the same tests.

What I changed

• detect MLX device architecture at draft load (applegpu_g13s on this machine)
• on applegpu_g13/g14, cast draft floating tensors to fp16 instead of leaving them in the default bf16 path
• prewarm verify kernels in fp16 when that cast is active

The patched server banner confirms it is live:

• Draft dtype: fp16 (applegpu_g13_g14_auto_fp16)

Results

Case	Config	Prompt tok	Gen tok	Prefill tok/s	Decode tok/s	Accept	TPC	Wall	Notes
35B-A3B before Tier 1	exact-body DFlash replay, draft-quant none, bf16 runtime	8392	324	413.70	19.34	0.660	2.95	37.04s	finish=stop
35B-A3B after Tier 1	exact-body DFlash replay, draft-quant none, fp16 runtime	8390	57	460.40	15.28	0.561	2.28	21.95s	finish=tool_calls; prefill improved, decode not apples to apples
27B dense before Tier 1	exact-body DFlash replay, W4/W4	8390	430	79.25	5.97	0.677	3.09	177.87s	baseline dense DFlash
27B dense after Tier 1	exact-body DFlash replay, W4 + fp16 runtime	8390	444	71.04	5.11	0.646	2.83	204.97s	worse than baseline
27B dense after Tier 1	exact-body DFlash replay, draft-quant none + fp16 runtime	8390	478	75.72	5.00	0.676	3.08	206.49s	most direct bf16->fp16 test
27B dense AR control	mlx_lm.generate, fixed 430, same prompt tokens	8392	430	83.10	16.69	-	-	126.94s	exact-body AR baseline

Read

• Tier 1 does help the MoE long-context prefill path on M1 Max: 413.7 -> 460.4 tok/s on the 35B-A3B exact-body replay
• But Tier 1 does not lift the long-context speculative decode floor in a useful way
• The cleanest signal is the dense 27B case:
  • with default W4 draft, decode goes down from 5.97 to 5.11 tok/s
  • with draft-quant none (the more direct bf16->fp16 test), decode is still only 5.00 tok/s
  • acceptance / tokens per cycle stay basically flat: 0.677 / 3.09 before vs 0.676 / 3.08 in the dense none + fp16 run
• So bf16 emulation is probably part of the story for prompt processing on M1/M2, but it does not look like the main cause of the real long-context speculative underperformance here

My takeaway: Tier 1 is real, but it does not fix the actual problem. This makes Tier 2 look more likely: SDPA ladder / verify path / long-context speculative overhead rather than bf16 alone.

[deepsweet]

Tier 1 is real, I've tried to convert draft models to float16 a while ago, and it's exactly about the PP boost.

• https://huggingface.co/deepsweet/Qwen3.6-35B-A3B-DFlash-FP16
• https://huggingface.co/deepsweet/Qwen3.6-27B-DFlash-FP16

I've also created a general benchmark to validate the bfloat16 emulation on M1/M2 the hard way.

A very similar struggle is happening with native MTP right now ml-explore/mlx-lm#990 (comment):

This caps the MTP benefit to ±1.1x or ±1.2x on M1/M2 (on the 27B at least), regardless of acceptance ratio. FP16 weights help from being in a degraded state to 1.1x, but then the bottleneck might be GPU scheduling at L=2.