[jvmenen]

Depends on #549. This branch is #549 plus one commit; only the last commit (perf(prefill): in-forward GDN boundaries for Qwen3.5/3.6 on the invariant lane) is new here. Once #549 is merged, the diff reduces to that commit.

Why

To reuse a prompt prefix later, a hybrid model needs a snapshot of its recurrent (GDN, Gated DeltaNet) layers at the restore position. Such a snapshot can only be taken at the end of a forward, so when a restore boundary falls inside the last prefill chunk, MTPLX cuts that chunk in two and runs a second forward. On Qwen3.6-35B-A3B that second forward reads all routed experts again: about 0.1 s per warm turn and per request with a session bank from 512 tokens. The prefill loops can already record boundaries inside one forward, but the mlx-lm Qwen3.5/3.6 models do not offer the hooks for it.

How it works, in plain terms

While a snapshot is requested, each GDN layer runs its own standard code twice, once up to the boundary and once after it, and saves its state in between. The attention and MoE layers still see the whole forward at once. Per GDN layer this is exactly the arithmetic of the two-forward route, so the result is bitwise the same, with one forward instead of two. When the model is loaded, MTPLX swaps each GDN layer for a version that can save its state halfway through a forward; the weights stay the same, only the code that runs them changes, and a model without GDN layers is left as loaded. This only happens together with the invariant lane from #549, because without that lane a wider forward changes the MoE routing. MTPLX_GDN_BOUNDARY_INFORWARD=0 brings back the two-forward route.

What changes

• New mtplx/gdn_inforward_boundaries.py publishes the two hooks the prefill loops look for, with the contract of qwen4_exp.py: boundary_capture_scope arms a capture of the requested positions, take_boundary_captures returns them once. A position counts only if every recurrent cache entry recorded it.
• _BoundaryCaptureMixin.__call__: while a capture is armed, each GDN layer runs its stock __call__ once per segment ([0,p), [p,S)) on its own cache and records cache.state in between. Per GDN layer that is exactly the ladder's arithmetic, conv tail and projections included, without copying mlx-lm code. Attention and MoE see the whole forward. With nothing armed it is the stock call.
• install_gdn_inforward_boundaries: class swap for qwen3_5.GatedDeltaNet and qwen3_next.Qwen3NextGatedDeltaNet; the parameter tree is unchanged.
• mtplx/runtime.py: installed only together with the lane, i.e. with MTPLX_BATCH_INVARIANT_PREFILL=1 on a model the lane admits; a refused model gets neither. The install report (server log and /health under degradation.batch_invariant_prefill) gains gdn_inforward_layers. The existing MTPLX_GDN_BOUNDARY_INFORWARD=0 restores the ladder.
• generation.py is unchanged: the loops and the restore path already handle in-forward captures.

Results

Apple M5 Pro 64 GB, macOS 26.6.2, MLX 0.32.2, Qwen3.6-35B-A3B MTPLX Optimized-Balance, profile turbo with production arguments, fan mode default, 2026-09-26, fork commit 198e7194 (this change on a branch that also carried other open PRs of mine, among them a session bank for /v1/completions and first-token logprobs). All variants with the lane and the session bank on; L = ladder (MTPLX_GDN_BOUNDARY_INFORWARD=0), H = hooks, in the order L1, H1, L2, H2, fresh server and SSD cache directory each. /health showed 30 GDN layers hooked and no missed captures.

Bitwise equality, H against L, both rounds:

• First-token requests: 240 of 240, probabilities and full top-20 equal.
• Prompt scoring: 243 of 243 prompts bitwise equal.
• NLL on two documents: all 5,590 and 4,618 positions equal.
• Greedy continuations 24 of 24, chat and agent transcripts at temperature 0 12 of 12.

Time and forwards (mean of both rounds):

	ladder	hooks	change
p50, prompts >= 512 tokens (51)	456 ms	366 ms	-20%
p90, all 240 prompts	459 ms	367 ms	-20%
prefill forwards per prompt >= 512 tokens (median)	2	1
p50, prompts < 512 tokens (189)	325 ms	322 ms	within noise
chat / agent conversation: prefill forwards	45 / 44	39 / 39
chat / agent: TTFT sum			-0.6% to -1.2%
decode	79.5 tok/s	79.7 tok/s	within noise

Without a session bank a completions request has no boundary in its last chunk, so there is nothing to gain there; chat and agent turns gain as shown.

Tests

New tests/test_gdn_inforward_boundaries.py (Metal): a tiny 4-bit Qwen3.5-MoE (3 GDN layers, 1 attention layer, routed and shared experts, 32 experts top-4) with the lane, through the real cold streaming loop and the warm suffix loop:

• Cold, 161 tokens: forwards [96, 64, 1] against 6 on the ladder. Warm, 60 cached tokens plus a 171-token suffix: [96, 74, 1].
• Both bitwise equal to the ladder: boundary positions, every leaf of every record (states and hidden), logits, hidden and cache.
• Without the lane, on a refused model (8 experts, top-2) or with MTPLX_GDN_BOUNDARY_INFORWARD=0: the ladder. Nothing armed equals the stock call; a partial capture is not a boundary; the runtime installs the hooks only on the lane.

Targeted tests (test_gdn_inforward_boundaries, test_batch_invariant_prefill, test_prompt_scoring_trunk_chunks) pass. Full suite with MTPLX_CONFIG=/nonexistent: 9,492 passed, 69 skipped, 1 failed: the same Laguna memory test as in #549 (test_laguna_s_2_1_ar_route_skips_qwen_performance_hooks, addressed in #535). python -m build and scripts/fresh_venv_smoke.sh pass. Ruff: new files clean, runtime.py unchanged against main.

🤖 Generated with Claude Code

[youssofal]

Not in 2.12.1, which was a stability release. I'll review this for the next release.