MLX Metal Kernels v0.20.0

MLX Metal Kernels v0.20.0

This release adds PR #32: the optimized prefill stack scaffold.

Highlights

• Added a new multi-layer prompt prefill path in ops/llama_prefill_ops.py
• Added layer-level and stack-level reference and composed prefill APIs
• Added token-id prefill helpers that turn prompt token sequences into embeddings, run stack prefill, and return logits plus filled KV-cache state
• Integrated explicit use_prefill=True support into the tiny synthetic generation pipeline
• Kept the older token-by-token decode-ingest path available through use_prefill=False
• Added a new prefill-then-decode demo and a standalone prefill benchmark
• Updated backend registry, unified benchmark runner, README, and roadmap/docs

What changed

New prefill scaffold

The new prefill module introduces:

• LlamaPrefillBackendConfig
• reference_prefill_backend_config()
• metal_prefill_backend_config()
• tiled_prefill_backend_config()
• fused_experimental_prefill_backend_config()
• reference_llama_layer_prefill()
• llama_layer_prefill()
• reference_llama_stack_prefill()
• llama_stack_prefill()
• prefill_token_ids()

This path processes prompt tokens in sequence form, fills one KV-cache per layer, and then allows decode to continue from the final prompt position.

Tiny pipeline integration

TinyGenerationPipeline now supports:

• use_prefill=True in config
• PrefillResult
• prefill_prompt(...)
• end-to-end prefill-then-decode generation for the synthetic tiny-model runtime

The pipeline still remains synthetic/random-weight only and does not claim production Llama or Mistral inference support.

Benchmarks and demos

Added:

• benchmarks/bench_llama_prefill_stack.py
• examples/prefill_then_decode_demo.py
• unified benchmark-runner coverage for llama_stack_prefill

Tests

Added focused coverage for:

• layer prefill reference vs optimized parity
• stack prefill reference vs optimized parity
• prefill-then-decode pipeline flow
• continued compatibility with the tiny generation demo/tests

Verification run for this release

Validated locally in the current environment with:

• python -m pytest tests/test_llama_layer_prefill.py tests/test_llama_stack_prefill.py tests/test_prefill_then_decode.py tests/test_tiny_generation_pipeline.py tests/test_full_tiny_generation_demo.py -q
• python examples/prefill_then_decode_demo.py
• python examples/full_tiny_generation_demo.py --prompt Hi --max-new-tokens 2 --greedy --backend-preset reference
• python benchmarks/bench_llama_prefill_stack.py --bits 4 --B 1 --S 8 --num-layers 1 --hidden-size 64 --intermediate-size 128 --num-heads 4 --num-kv-heads 2 --head-dim 16 --MAX_S 16 --dtype float16 --backend-preset reference --iters 2 --with-lm-head

Known limitations

• The release is still synthetic/random-weight only
• Contiguous KV-cache is the implemented prefill path in this scaffold
• Continuation prefill with start_position > 0 is intentionally conservative and may raise NotImplementedError
• Paged prefill is future work
• Apple Silicon / Metal execution was not verified in this Windows environment
• Full python -m pytest tests -q remains blocked here by missing mlx.core

Upgrade notes

If you were using the tiny generation scaffold already, the public path remains backward compatible:

• existing decode-ingest generation still works
• prefill can now be enabled explicitly through TinyGenerationPipelineConfig(use_prefill=True)

Next roadmap items

• tokenizer/checkpoint package alignment
• local real-model smoke test
• quantized package tensor-data writer
• paged prefill