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.
* renderers: add gemma4-nemotron renderer
Add a Gemma 4 renderer variant that renders assistant content before
its tool calls rather than after the tool responses. Models opt in with
RENDERER gemma4-nemotron; gemma4 auto-detection is unchanged.
Adds the corresponding chat template as a test fixture for Jinja parity
checks.
* mlx: support Gemma 4 targets with DFlash drafts
Load a draft-owned output head when its config requires one, while preserving the target-head path used by existing DFlash models.
Expose Gemma 4 hidden-state taps and raw logits for speculative decoding. Report safetensors target precision independently of quantized draft weights.
model/base held the Model interface and the architecture registry while
model held weight loading and quant parameters, and every architecture
imported both. create/client was the CLI side of safetensors imports,
with the create command as its only caller, a duplicate of the command's
adapter error, and a name that read like a second API client. The
safetensors show helpers had their own package under x/ although it
already declared package server.
base merges into model, so base.Model and base.Register become
model.Model and model.Register; nothing in the two overlapped. The create
client's two files and their tests join package cmd, and the five names
the command called are no longer exported. The show helpers join the
server package, and the three entry points routes.go calls become
unexported like the helpers around them.
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