[davidtai]

Qwen3.8 Flash-Next: two verify-step memory fixes so 261,120-token prompts decode

This pull request adds two serving optimizations to the Qwen3.8 Flash-Next path. The base is upstream MTPLX 2.11.2. At the pack's advertised 261,120-token context, mtplx serve reads the whole prompt and then runs out of GPU memory in the first decode step, so the request fails after a few tokens. Both optimizations remove a per-layer allocation the speculative verify makes across the 12 full-attention layers, and both are needed: the verify's KV-cache write is what still overflows once the attention transient is gone.

The two are ordered by what they remove. Optimization 1 writes the verify's KV cache in place instead of reallocating the whole buffer, which is the allocation that decides the 261,120-token cell. Optimization 2 keeps the verify attention on MLX's fused kernel, which removes an attention transient that scales with context; on its own it does not make the cell fit, and Section 3 gives its measured 261,120-token receipt.

The comparison has two arms: release 2.11.2 alone (A) and release plus both optimizations (F, this pull request). The headline is F against A at 261,120 tokens, where A exceeds the memory knob and F does not.

Headline figures:

• 261,120 tokens, cold prefill: release 2.11.2 exceeds the 100 GiB memory knob and returns insufficient_memory in the first decode step; this pull request completes the request, fits ×3, at 62.84 tok/s decode (56.82-62.84) and 1082.9 tok/s prefill, peak 100.82 GB against the 107.374 GB knob.
• 16,384 tokens, cold prefill: decode is unchanged. Release reaches 71.97 tok/s fastest (71.58 mean over its ABAB windows) and this pull request reaches 71.61 tok/s; both optimizations engage only in the verify step, so neither moves the token rate.
• Both optimizations are independent of pull requests [FEAT] 4 unimplemented optimizations from PR391 plus 2 new ones - six serving optimizations for Qwen3.8 Flash-Next on MTPLX 2.11.2 (four PR-391 remainder optimizations + two exact decode optimizations) #475 and [FEAT] Typical Acceptance Verifier from Medusa2 - perf(qwen4): typical acceptance behind --typical-threshold (default off) #478. All three are based on upstream main, and this one fixes the 261,120-token cell that both of those report as exceeding the knob.

Terms used in this document:

• decode: the token-by-token generation phase, after the prompt is read.
• prefill: reading the prompt into the KV cache before the first output token.
• verify: the speculative verify step inside decode; the fixed four-row width the compiled Flash-Next verify uses.
• optimization: one change an operator can switch off on its own.
• fixed-M4 verify: the four-row speculative verify the served Flash-Next path runs; it holds its KV in TensorOffsetKVCache rather than the stock in-place KVCache.
• QSA: the sparse-attention path, which selects the earlier tokens each layer attends to.
• GQA: grouped-query attention; here 24 query heads over 2 key/value heads, so 12 query heads per key/value head.
• full-attention layer: one of the 12 of 48 layers that keep dense KV over the whole context.
• rounding-class: the output differs from the replaced path only by floating-point rounding.
• tok/s: tokens per second.

---

1. Benchmark results

Two arms ran on one machine. Each arm received the same request body per cell. Section 6 gives the settings.

1.1 Peak memory by context

Peak memory by context size, release against this pull request. The 100 GiB cap holds on every arm that reaches a cell; release exceeds it at 261,120 tokens and this pull request does not.

Context	release 2.11.2 (A)	this PR (F)
16,384	93.41	93.41
261,120	exceeds the 100 GiB knob (Metal insufficient_memory)	100.82

1.2 261,120-token result

Decode tok/s at 261,120 tokens. Release does not reach a decode rate because it fails in the first decode step; this pull request completes the request.

Release 2.11.2 does not complete this cell. Its receipt shows the prompt fully read and only a few tokens emitted before the failure: new_prefill_tokens 261120, time to first token 243.5 s, completion_tokens 12, finish_reason error, sampled peak 100.82 GB, with [METAL] Command buffer execution failed: Insufficient Memory. Section 2 traces the cause.

This pull request completes the request in three cold windows:

Window	Status	Prefill tok/s	Decode tok/s	Time to first token (s)	Wall (s)	Completion tokens	Peak memory (GB)
w1	fits	1081.6	59.62	242.631	254.03	461	100.82
w2	fits	1081.5	56.82	242.614	264.33	1024	100.82
w3	fits	1082.9	62.84	242.290	259.43	845	100.82

The summary reports the fastest window with the min-max across the three: prefill 1082.9 tok/s (1081.5-1082.9), decode 62.84 tok/s (56.82-62.84), time to first token 242.290 s (242.290-242.631), wall 254.03 s (254.03-264.33), peak 100.82 GB (100.82-100.82). Peak stays under the 107.374 GB knob on every window, with about 6.5 GB of headroom. Completion tokens differ by seed (461-1024) because each seed stops on its own; every window reached a normal stop or the 1,024-token cap, never an error.

1.3 16,384-token no-regression check

Both optimizations engage in the verify step, which runs at every context, but 16,384 tokens never overflowed, so this cell measures that the token rate is unchanged. This pull request's arm F is release 2.11.2 plus both optimizations, so the comparison is F against release A.

Metric	release 2.11.2 (A)	this PR (F)
Decode tok/s (fastest, min-max)	71.97	71.61 (71.24-71.61)
Decode tok/s (mean of windows)	71.58	71.47
Prefill tok/s	1381.5	1377.2 (1343.5-1377.2)
Time to first token (s)	12.056	12.112 (12.112-12.390)
Peak memory (GB)	93.41	93.41 (89.01-93.41)
Completion tokens	1024	1024
Output vs release	reference	byte-identical, all 3 seeds

For context, pull request #475's six-optimization arm reaches 82.84 tok/s fastest (76.91 mean) at this cell on the same instrument. That arm is a separate, larger change; this pull request is the memory fix alone and does not aim to move the 16K token rate.

---

2. The 261,120-token overflow

The prefill is not what overflows. All four arms of the earlier 261,120-token battery, release 2.11.2 (A), release plus #475 (C), and the two #478 arms (D and E), show the same shape: the full prompt is read and tokens are already emitted before the failure.

Arm	new_prefill_tokens	time to first token (s)	completion tokens	finish reason	sampled peak (GB)	process footprint (GB)
A release	261120	243.49	12	error	100.82	102.83
C plus #475	261120	241.88	3	error	100.79	102.87
D plus #478 at 0.20	261120	241.66	7	error	100.79	102.82
E plus #478 at 0.09	261120	240.46	10	error	100.79	105.32

The error field is identical on all four:

server error [insufficient_memory] insufficient memory: the request exceeded available
GPU memory ([METAL] Command buffer execution failed: Insufficient Memory
(00000008:kIOGPUCommandBufferCallbackErrorOutOfMemory).). The engine shed its caches and
stays up; this request failed.

The overflow is in the first decode step, which is the fixed four-row speculative verify. Two separate per-layer allocations sit in that step, both multiplied by the 12 of 48 layers that keep dense KV over the whole context (full_attention_interval 4). The model has 24 query heads over 2 key/value heads, so 12 query heads per key/value head, at head_dim 256.

The attention transient. A multi-row verify has a query length of 3 to 6 rows, and 3 rows times the group factor of 12 is already past the 32-row limit of MLX's fused vector attention kernel. MLX declines the fused call, and the unfused route runs instead, materializing a [24, S, 261120] score tensor per full-attention layer and expanding keys and values per query head. That transient scales with context.

The KV-cache write, which is the binding one. The fixed-M4 verify holds its KV in TensorOffsetKVCache, whose update_and_fetch wrote the new rows with the functional mx.slice_update (mtplx/graphbank.py:499-522 before the fix, the write calls at :514 for keys and :520 for values). That op returns a new array: it reallocates the whole [1, 2, capacity, 256] bf16 buffer for keys and again for values, about 267 MB each at the 262K capacity, so about 6.4 GB across the 12 full-attention layers. The stock KVCache on the plain autoregressive path writes in place and allocates nothing.

Both land on a steady state that is already about 100.8 GB, against a 107.4 GB limit, so the single decode command buffer exceeds the knob. At 131,072 tokens the same terms are about half and the request fits; at 261,120 tokens they do not.

Evidence chain

Four measurements identify the KV write, in the order they were taken.

W5, native MTP off. Decode becomes plain autoregressive at query length 1: the fused kernel is eligible so no score plane is built, and the stock in-place KVCache runs instead of TensorOffsetKVCache. That configuration fits the same 261,120-token prompt at 94.19 GB and decodes 547 tokens to a normal stop, time to first token 223.6 s. With MTP on and the verify running, the peak is 100.82 GB. The 6.4 GB gap is what the verify adds, and turning MTP off removes both candidate terms at once, so on its own it does not say which. Receipt: battery475/receipts/arm-C-w5-nomtp/.

Head chunk alone. Removing only the attention transient left the peak unchanged at 100.82 GB and the request still failed. That isolates the remaining 6.4 GB to something other than attention. Section 3, "What did not fix it".

The allocation probe. A direct probe of TensorOffsetKVCache at this geometry (context 261,120, capacity 261,128, 12 layers, 5 verify rows) measures one KV buffer at 0.27 GB, resident KV across keys, values and the 12 full-attention layers at 6.42 GB, and, with the in-place write, an update transient above resident of 0.00 GB. Unfixed, that transient is the full 6.4 GB again, because mx.slice_update returns a new buffer per tensor per layer. Log: battery475/kvprobe/kvprobe-20260908T005614Z.log.

The fix arm. With the in-place write, the same prompt fits on all three cold seeds at a 100.82 GB peak, 6.5 GB under the 107.374 GB knob, with MTP on and the verify running. Section 1.2.

The probe's 6.42 GB resident and 0.00 GB transient, against the 6.4 GB gap W5 exposed, is the arithmetic that closes the chain.

The served log records the refusal verbatim, one line per shape class:

[scaled_dot_product_attention] force_fused=True but no fused kernel is available: the
vector attention kernel requires the query length times the GQA factor to be at most 32;
got query length 3 and GQA factor 12.

The refused shape classes observed at this cell are query length 3, 4, 5 and 6. The wide 256-row prefill chunks still fuse, so prefill is unaffected; the served log shows fused SDPA for shape class [causal-mask q_len 256 x GQA 12 at head_dim 256 bfloat16] in the same run. The QSA sparse layers avoid the transient through their own bounded row selection; the 12 dense full-attention layers cannot, and the fixed four-row verify path additionally turns off the score tiler and the rows-gather route, which is why lowering the allocator cache limit, turning the session bank off, and the earlier score-tiler probe were each no-ops at this cell. The full analysis trail is in over100-reports/pr-bodies/261k-correction-report.md and over100-reports/oom255k/report.md.

---

3. What the optimizations do

#	Optimization	Phase	Exactness	Key
1	In-place verify KV write	decode	exact (byte-identical write)	none; always on for a concrete offset
2	Verify SDPA head chunk	decode	rounding against the unfused route	MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK

1. In-place verify KV write (the fix)

• Problem. TensorOffsetKVCache.update_and_fetch wrote the verify's new KV rows with the functional mx.slice_update, which returns a new array rather than writing into the existing one. At the 262K capacity that reallocates a [1, 2, capacity, 256] bf16 buffer, about 267 MB, once for keys and once for values, in each of the 12 full-attention layers: about 6.4 GB inside one verify command buffer. The stock KVCache used by plain autoregressive decode writes in place and allocates nothing, which is why the same prompt fits with native MTP off.
• Change. update_and_fetch, and the trim rollback restore, take an in-place path when the cache offset is a concrete host value: the S new rows are written with a slice assignment into the buffer that already exists, after materializing an independent copy of the pre-write rows for rollback. The offset is concrete on the eager verify path, which is the path the served 261,120-token request takes, because the bank demotes long generations to the eager verify. When the offset is a tracer, that is, inside an mx.compile trace, the compile-visible functional slice_update path is left exactly as it was, so the compiled replay's buffer-donation contract is untouched.
• Effect. The 6.4 GB of per-layer reallocation in the verify step goes to zero, which is the term that decides the 261,120-token cell.
• Exactness. Exact, not rounding-class. The in-place write stores the same bytes slice_update would have stored, verified on random tensors, and the rollback snapshot is materialized independently before the write, so trim restores exactly the pre-write rows.
• Files. mtplx/graphbank.py.
• Switch. No key. The in-place path is taken whenever the offset is a concrete host value; the compiled path is unchanged, so there is nothing to opt out of.

2. Verify SDPA head chunk (MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK)

• Problem. In the fixed four-row speculative verify, the 12 full-attention layers attend over the whole KV cache with a query length of 3 to 8 rows. Query length times the group factor of 12 exceeds MLX's fused vector attention limit of 32, so SDPA falls to an unfused route that materializes a [24, query length, context] score tensor per layer and expands keys and values per query head. That transient scales with context.
• Change. The optimization splits the query heads of that attention into chunks small enough that each fused call satisfies query length times heads-per-chunk of at most 32 (for example query length 6 gives chunks of 5 heads, query length 4 gives chunks of 8 heads). Each chunk is paired with its own single key/value head, so keys and values are never expanded per query head, and the chunk outputs concatenate along the head axis. Every dispatch stays on the fused kernel, so no [24, query length, context] tensor is built and the working set is one chunk regardless of context. Wide prefill chunks, at query length above 32, keep their single fused call. Stock MLX, no native extension.
• Effect. The context-scaling attention transient in the 12 dense layers is removed. On its own that is not enough to fit the 261,120-token cell, which the KV reallocation of optimization 1 still overflows; see "What did not fix it" below. It is exact and it costs nothing, and it keeps the verify on the bounded fused kernel so the attention term does not return as context grows.
• Exactness. The heads in one chunk all belong to the same key/value head, so each chunk's attention is exactly the full grouped-query attention's arithmetic for those heads on the same fused kernel. The chunked result is bit-identical to an unchunked fused call, verified below 1e-6 on fp32 random tensors against a grouped-query-expanded reference. Against the unfused route it replaces, it is rounding-class, differing only by floating-point accumulation order, the same class as the prefill mask-fuse optimization.
• Files. mtplx/models/qwen4_exp.py.
• Switch. Default on for a served Flash-Next pack; it engages only in the band where query length times the group factor exceeds 32 and query length is at most 32. MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK=0 turns it off. The key is read at use, so the server's auto-arm reaches it.

Optimization 2 applies at the single dense-fallback attention call, so it covers the fixed four-row verify, the ordinary verify, and any small prefill tail chunk in the same band.

What did not fix it

Optimization 2 was measured alone at 261,120 tokens before optimization 1 existed, and it did not fit. The head chunk engaged, GET /health reporting verify_sdpa_head_chunk {engaged: true, q_len: 5, heads_per_chunk: 6, chunks: 4, n_kv_heads: 2}, so the [24, S, 261120] score plane was not built; the request still failed in the same place, at the same peak:

arm	new_prefill_tokens	time to first token (s)	completion tokens	finish reason	sampled peak (GB)
A release 2.11.2	261120	243.49	12	error	100.82
head chunk only	261120	242.44	12	error	100.82

Receipt: battery475/receipts/arm-F-headchunk/arm-F-headchunk-261120-s20260829-*. A peak unchanged to two decimals, with the attention transient provably gone, is what pointed at a second per-layer allocation of about the same size and led to the KV-cache write.

The same arm at 16,384 tokens is a clean no-regression: three cold windows, decode 71.11 to 71.56 tok/s, and the output is byte-identical to release on all three seeds, text_sha256 matching arm A seed for seed (ee9993ec..., eec89cc6..., 1ac58119...) with the same completion-token counts. Receipt: battery475/receipts/arm-F-headchunk/arm-F-headchunk-16384-*.

---

4. Independence from pull requests #475 and #478

This pull request is based on upstream main (MTPLX 2.11.2) and touches one function in mtplx/models/qwen4_exp.py and one in mtplx/graphbank.py. It does not depend on #475 (the PR 391 remainder and two exact decode optimizations) or #478 (typical acceptance), which are also based on upstream main. It fixes the 261,120-token cell that both of those pull requests report as exceeding the memory knob. The three can land in any order; if #475 or #478 land first, these changes apply on top with no conflict, since they touch one attention call and one KV-cache write that the others do not.

---

5. Quality and tests

Optimization 1 is exact: the in-place write stores the bytes slice_update would have stored. Optimization 2 is rounding-class against the unfused path it replaces and bit-identical to an unchunked fused call. Correctness rests on the equivalence tests below plus the repo's code-eval gate.

CPU tests in tests/test_tensoroffset_kv_inplace.py:

• the eager update is in place, asserting no full-buffer allocation;
• it writes the correct rows and advances the offset;
• the rollback snapshot is independent of the in-place write, and trim restores the pre-write rows;
• the eager write matches the functional mx.slice_update numerically;
• the concrete-offset probe distinguishes a tracer from a host value, so the compiled path keeps the functional write.

CPU tests in tests/test_qwen4_verify_sdpa_head_chunk.py:

• the key resolves at use and defaults on, replayed in the served import order so an import-frozen reader would be caught;
• the chunk plan for query length 1 through 8 (and the prefill regime) at group factor 12, with the invariant that every chunked call keeps query length times heads-per-chunk at most 32 and covers all heads;
• numerical equivalence, that the chunked fused result matches an unchunked grouped-query-expanded fused reference within 1e-6 for query length 3 through 8;
• a dispatch spy that asserts every fused call stays within the 32-row bound and pairs one key/value head, so no [24, query length, context] tensor is ever built and keys and values are never expanded;
• the opt-out path takes the single unchunked call, and query length 1 decode is never chunked;
• the engagement receipt for /health.

The existing tests/test_health_qwen4_install_reports.py still passes with the receipt added.

Optimization 2's engagement is observable without the serve log. On first engagement the server prints [mtplx] verify SDPA head-chunked: q_len=S heads_per_chunk=H chunks=N, and GET /health carries qwen4_install_reports.verify_sdpa_head_chunk with the query length, heads-per-chunk, chunk count and key/value head count once it has fired. Optimization 1 has no receipt of its own; it is observable as the peak-memory drop at long context and is covered by its CPU tests.

---

6. Benchmark method

Setting	Value
Model	Youssofal/Qwen3.8-Flash-Next-MTPLX-Optimized-Speed, revision 29ba90f82124961d0d902a9ea9bbb1034972af2f
Arms	release MTPLX 2.11.2 (A); release plus both optimizations (F, this pull request). A head-chunk-only arm was measured separately and is reported in Section 3 under "What did not fix it"
Engine build	MTPLX 2.11.2 on MLX 0.32.2
Profile	cli-resolved (Turbo), the profile mtplx serve selects for this pack; a fresh SSD session cache directory per window
Prompt sizes	16,384 and 261,120 tokens
Output length	1,024 tokens maximum per request
Sampler	temperature 1, top-p 0.95, top-k 20
Reasoning effort	xhigh
MTP depth	native, depth 3
Seeds	20260829, 20260830, 20260831
261,120 protocol	three cold windows on arm F; each must complete under the knob
16,384 protocol	interleaved windows; each arm's value is its fastest window
Cell statistic	throughput cells report the fastest of the windows; time to first token and wall report the fastest (lowest); peak memory reports the highest
n-gram table	pre-warmed before each cell
Memory cap	100 GiB on every arm
Prefill state	cold; cross-request prefix restore off, so each window prefills fresh; new prefill tokens equal prompt tokens
Thermal state	fans at maximum, a 40 degree Celsius gate before every cell

---

7. How to run and how to disable

Serve the pack; the server arms both optimizations by default:

mtplx serve \
  --model ~/.mtplx/models/Youssofal--Qwen3.8-Flash-Next-MTPLX-Optimized-Speed \
  --model-id mtplx-flash-next-optimized-speed

Read GET /health and confirm qwen4_install_reports.verify_sdpa_head_chunk once a long-context request has decoded, or watch the serve log for [mtplx] verify SDPA head-chunked:. The in-place KV write has no key and no receipt; it shows up as the long-context peak-memory drop.

Turn optimization 2 off with its key set to 0:

MTPLX_QWEN4_VERIFY_SDPA_HEAD_CHUNK=0 mtplx serve --model ...

---

8. File map and provenance

Area	Files
Runtime	mtplx/graphbank.py (optimization 1), mtplx/models/qwen4_exp.py (optimization 2)
Server	mtplx/server/openai.py (the /health install receipt)
Tests	tests/test_tensoroffset_kv_inplace.py, tests/test_qwen4_verify_sdpa_head_chunk.py
Documentation	docs/perf/qwen38-longctx-prefill-memory.md; analysis trail over100-reports/oom255k/report.md and over100-reports/pr-bodies/261k-correction-report.md

• Base: upstream main 21be78b3 (MTPLX v2.11.2).
• Branch: perf/qwen38-longctx-prefill-memory-main @ 9f306e01 = upstream main, then two code commits and one docs commit. The code commits:
  • e0ebdc41 "Head-chunk the small-q_len verify SDPA to keep long-context decode fused" (optimization 2);
  • fc0d2a29 "Write the fixed-M4 verify KV in place to fit long-context decode" (optimization 1, the fix).
• Diff against base: 6 files, mtplx/graphbank.py +75/-16 region, mtplx/models/qwen4_exp.py +144, mtplx/server/openai.py +12, two test files (+165, +239) and docs/perf/qwen38-longctx-prefill-memory.md.
• Results docs: 9f306e01 "docs(perf): 261,120-token decode fix results, evidence chain, and receipts index" adds docs/perf/qwen38-256k-fix/ (README plus the two charts this body renders). Full SHA 9f306e015073a35e0797650d7796d27d5749fe8b.

[davidtai]

This would blow the fixed budget by about 6.6gb due to the copy. I have a hard limit on my mac so it would risk OOM and terminate.

[youssofal]

Agreed, and that is why the in-place write is the right shape. Copying the verify KV at that context costs about what you say, and we run under the same hard budget here, so a copy was never an option.

Your three commits are already on the 2.11.3 branch with your authorship intact. I rebased them onto the release line this morning, so the hashes differ from this branch and GitHub will not mark the pull request merged on its own; I will close it when the release goes out. Your two test files pass there alongside our full suite, and on our 128 GB machine the head-chunk path shows as engaged in /health on the running daemon. Thank you for this one.

[youssofal]

Merged for 2.11.3, your commits 69f222f and 57fa917 with their authorship kept (merge f5ae03d), thank you. 261,120-token prompts decode on Flash-Next now, and the release notes credit you. Closing here because the merge went in by patch rather than through this button. Release: https://github.com/youssofal/MTPLX/releases/tag/v2.11.3