MLXLLM: add Recursive Flint (flint-int4) model

Merged
SI
10d ago
+882-0

Adds Recursive Flint (model_type: flint) so RecursiveSI/flint-int4 loads and generates. It's a port of mlx_lm/models/flint.py from ml-explore/mlx-lm#1998, which is still open upstream.

Flint is a 40B MoE with about 2.5B active per token. It has 49 layers: 2 dense, then 47 with 256 experts routed top-6 by a sigmoid router, plus a shared expert. 35 sliding-window layers (window 513) use RoPE, and 14 full-attention layers have no positional encoding. Every layer has QK RMSNorm and learned attention sinks.

The INT4 checkpoint needed two loader fixes, both in sanitize. It stores packed codes as uint8, so they get reinterpreted as uint32. Its embeddings are 8-bit even though the config only declares 4-bit, so sanitize reads the real width off the tensor shapes and quantizes that layer itself; the loader skips modules that are already quantized. The expert top-k is written inline, because ra/main has no router helper.

A quantized KV cache is refused with a clear error, because quantizedScaledDotProductAttention has no sinks argument and would silently give wrong output.

Testing:

  • FlintTests: 24 hermetic tests on a tiny random config, giving 100% line coverage of Flint.swift. They cover config validation, cached vs uncached decode past the window, sinks, tied embeddings, and the checkpoint repack.
  • FlintRealModelTests: opt-in through MLX_FLINT_MODEL. It loads the real flint-int4 weights and checks that greedy decoding of def add(a, b):\n return gives a + b. It passes on an M-series Mac with 48 GB.
TEST_RUNNER_MLX_FLINT_MODEL=~/Downloads/flint-int4 xcodebuild test -scheme mlx-swift-lm-Package \
  -destination 'platform=macOS' -skipPackagePluginValidation \
  -only-testing:MLXLMTests/FlintTests -only-testing:MLXLMTests/FlintRealModelTests

Only the Flint tests were run; the rest of the suite wasn't. Distributed sharding isn't ported.