Skip to content

TensorFold 0.3.4.1: prompt processing back to full speed

Choose a tag to compare

View all tags
@ashhart ashhart released this 27 Sep 08:47
v0.3.4.1
bb4b4a3

A fix for 0.3.4: prompt processing is back to MLX's speed on every model, and a resumed conversation still matches the same conversation fed fresh, byte for byte.

What was wrong

0.3.4 processed Nemotron's and Flash Next's prompts through their row-exact decode kernels, so that a conversation resumed from a cached prefix got exactly the bits of a fresh one. That made prompt processing several times slower than MLX's own: about 550 tok/s for Nemotron on an M5 Max, where mlx_lm runs about 2,500-3,000, and about 300 for Flash Next on an M3 Ultra, where MLX runs 760.

What changed

  • Every model prefills through MLX's own forward, in chunks on a fixed 2,048-token grid from the start of the prompt. Prompt caches are kept only at grid points, and a reply is prefilled again from the last grid point on the next turn. A chunk's bits then never depend on where a conversation was resumed, so resumed conversations still equal fresh ones.
  • Flash Next queues at most two layers ahead of the GPU while it prefills, so its peak memory stays bounded at any context length.
  • The server no longer holds the first request while it warms saved system blocks at startup; the warm runs in the background.
  • Prompt caches saved by 0.3.4 are not reused, because their key now names the prefill mode, so the first prompt after upgrading is prefilled from scratch.

Measured on this tree

Cold prompts, tok/s, through tensorfold serve:

Model Machine Prompt 0.3.4 0.3.4.1 MLX
Nemotron 3.5 Lightning M5 Max 8k ~550 3,400 2,540 (mlx_lm, same session)
32k ~540 2,456 2,448 (mlx_lm, same session)
Qwen3.8-27B M5 Max 2k 542 867 812 (mlx_lm)
32k 510 560 562 (mlx_lm)
64k 441 465 467 (mlx_lm)
Qwen3.8 Flash Next M3 Ultra 2k 305 934
16k ~290 820 759 (MLX, at 22.8k)
64k ~280 550
196k ~255 315

Flash Next's peak memory stays within 20 GB of its weights up to a 196k-token prompt (134 GB peak on the M3 Ultra). Its prompt speed still falls with context length; a sparse prefill kernel that reads only each row's selected keys is next.

Checks on this tree:

  • Every drafted reply equals the same request sent with "draft": false, 9 of 9 on each model.
  • Resumed multi-turn conversations equal the same conversations fed fresh, including a three-turn Flash Next chat resumed at nine grid points of a 19k-token first message, with thinking on and off.
  • Decode speed is unchanged: Nemotron on an M5 Max within noise of 0.3.4 over 15 seeds a cell, and Flash Next on an M3 Ultra at 135.5 / 111.1 / 142.3 / 122.2 against 134.1 / 111.8 / 137.7 / 117.0 (code sampled / chat sampled / code greedy / chat greedy).
  • The test suite passes (370 tests).

Install

tensorfold update installs this release. Or:

pip install --upgrade git+https://github.com/ashhart/TensorFold.git@v0.3.4.1
Assets 2
Loading