MLX's argpartition sorts the whole row, and a row longer than one block
of the GPU sort takes a multi-pass merge sort. TopKIndices sorts each
2048-value chunk in a single block, then only the chunks' top k, and
returns the same indices in the same order.
The sampler's top_k and top_logprobs and the DFlash2 candidate selector
use it in place of a full-vocabulary argpartition.
DFlash2 adds block-local dynamic convolutions and a candidate selector
to DFlash. The selector proposes each token conditioned on its
predecessor and returns each position's candidate distribution, which
verification uses to keep the output exact. Like the reference, the
draft samples with the request's temperature only.
A DFlash2 session drafts seven tokens every round on the existing
snapshot verify path instead of following the depth controller, whose
shallower picks lose more accepted tokens than they save in verify time.
qwen3.5-family targets now expose the layer taps, token embeddings and
raw logits a DFlash draft needs.
The MLX runner is the only Go inference runner left and is no longer
experimental, so its packages leave x/. The bindings become a top-level
mlx package beside the carried patches in mlx/compat, mirroring how
llama/ holds the llama.cpp integration, and the runner becomes mlxrunner
with the architectures nested under the package they implement.
Subpackages move with their parent unless listed.
x/mlxrunner/mlx mlx
x/internal/mlxthread mlx/mlxthread
x/internal/mlxthreadtest mlx/mlxthread/mlxthreadtest
x/internal/mlxtest mlx/mlxtest
x/quant mlx/quant
mlx/compat/*.patch mlx/compat/mlx-c (MLX patches go in mlx/compat/mlx)
x/mlxrunner mlxrunner
x/models/nn mlxrunner/nn
x/models/<arch> mlxrunner/model/<arch>
x/mlxrunner/imports.go mlxrunner/model/architectures (new package)
x/create create
x/safetensors fs/safetensors
x/tokenizer mlxrunner/tokenizer
Every package keeps its name, so the Go changes are the import path
rewrites the moves force, and the CMake, Dockerfile, CI cache keys, drift
check and Darwin payload script follow the new paths. Four edits are not
paths: the runner's blank architecture imports become the package
mlxrunner/model/architectures, so the list to extend for a new model sits
beside the architecture directories; a depguard rule keeps the two test
harnesses out of non-test code, as the x/internal placement used to; the
CI change filter's two entries for the long-deleted x/imagegen/mlx now
name the bindings' CMake project and the carried patches, so a change to
either builds the payload; and the tokenizer parity test reads its
fixtures from its own testdata instead of walking out of x/.
x/server and x/imagegen/manifest stay for the next two commits.
Daniel HiltgenandGitHub
Jesse Gross