Many small perf improvements, add Gemma 4 family Qwen 3.5+ family, vision and audio - #356
Open
flukejones wants to merge 41 commits into
Open
Many small perf improvements, add Gemma 4 family Qwen 3.5+ family, vision and audio#356flukejones wants to merge 41 commits into
flukejones wants to merge 41 commits into
Conversation
Zero-byte file carried in from upstream f4aa309; never had content.
Review checks that compile clean but fail or silently corrupt at runtime on the Metal stream — the class of bug clippy can't catch. - MLX state/threading/Metal: no thread_local for compiled-graph or Array caches (mlx-c v0.31 SIGSEGVs on the destructor/stream race), kernel-name caching, grid semantics, f64 reject, bool-mask sentinels. - FFI: mlx-sys is the only direct mlx-c caller; Drop-once on handles; single-source version pin. - Array clones as per-token FFI roundtrips; borrow-not-clone; async_eval decode scheduling; no item() barrier in the hot loop. - Dtype strictness: silent bf16->f32 promotion poisons the graph; cache dtype-bound scalars. - Fixed-N allocation, decode-only bench methodology, error/type/import conventions.
Harness: - criterion harness; short + long_prompt (T=1024) per model - llama 1B + qwen3 1.7B at bf16/q8/q4 - MLX_LM_BENCH_SET=full adds llama 3B + qwen3 0.6B at same quants - MLX_LM_BENCH_ONLY substring filter on per-cell group prefix - checkpoint cache + auto-download via `hf` CLI - MLX_LM_BENCH_NO_DOWNLOAD, MLX_LM_BENCH_CACHE, MLX_LM_BENCH_SET - BENCHMARK.md methodology + run instructions Shared decode step (so bench can't drift from production): - models::decode_step — one per-token unit (reshape, forward, slice last position, sample). Both Generate::next() and the bench drive it. - decode is synchronous here (no async pipelining yet); the bench times Generate::next() one token at a time, eval-fenced — the real cost. - unify the duplicated per-model ModelInput into nn::ModelInput. - fix qwen3 decode: it sampled un-sliced [B,1,vocab] logits (NewAxis then stacked each step); now slices last position like llama. Also hoists pre-existing crate::utils::scaled_dot_product_attention qualifiers in both models to top-level use (check-paths). Pre-runtime scaffolding: models::decode_step + the per-model Generate iterators benched here are the only decode path until the LanguageModel/ ModelContext runtime lands, which replaces both and rewrites this bench onto the production ModelContext path. The throwaway is intentional — the runtime can't precede the models it abstracts.
Hand-picked from clippy pedantic; pass-by-value on public
constructors deferred to perf-scope commit where the API break
has a measured justification.
mlx-rs:
- iter().cloned() -> copied() on i32 slices (linear, indexing)
- iter().for_each() -> for loop (module trait)
- iter() removed from for-loop heads (modules, optimizers, io)
- match Some/None -> let-else (utils::mlx_closure_payload)
- match self.is_frozen { true/false } -> if/else (param)
- match self.nesterov { true/false } -> if/else (sgd)
- match align_corners -> if/else (upsample)
- panic-in-if -> assert! (indexing ellipsis check)
- needless borrow on randint scratch (arithmetic test)
- allow(match_same_arms) on Dtype::promote_with table:
per-pair enumeration is intentional documentation
mlx-lm:
- match (k,v) -> let-else for quantized/non-quantized pair
(utils::scaled_dot_product_attention)
mlx-lm-utils:
- iter().map(..).flatten() -> flat_map (tokenizer test)
mlx-tests:
- allow(dead_code) on QuantizableExample (derive-only example)
Restores cargo clippy --workspace --all-targets -- -D warnings.
- bump mlx-c submodule a1290d -> fba4470 (v0.5.0 -> v0.31.2) - workspace 0.25.3 -> 0.26.0 - Stream::sync, Array::wait + is_available - nn::RmsNorm weight Option (no-scale variant) - transforms::compile rewrite: TypeId cache, GAT CallMut, per-arity shape markers; closures +Send - fast::MetalKernel + MetalKernelConfig - random default state thread-local (v0.31 TLS hazard) - memory module - fft / quantization ops adapt to v0.31 signatures - mlx-rs/benches/compile_overhead - nn::Rope + mlx-lm utils::rope: pass 4-D `[B,N,T,D]` direct; prior reshape zeroed all but head-0 on decode - nn::Alibi: bias-matrix cache moved from a thread_local to a struct field — v0.31 SIGSEGVs when a thread_local Array cache destructor races the GPU stream at thread exit. Alibi gains the cache field + Default; matrix() takes &self. - mlx-lm: quantization module + qwen3/llama loader rewrites `<prefix>.weight` -> `<prefix>.inner.weight` for keys with a `.scales` sibling (q4/q8 against v0.31 QuantizedLinear) - qwen3 loader handles single-file + sharded checkpoints
Squashed port of 12 fork commits targeting upstream-resident files.
perf:
- defer eager FFI calls in unwrap_or hot paths: 27 sites switched to
unwrap_or_else so mlx_array_new only fires on None
(fast::*, ops::quantization::*, ops::arithmetic::block_masked_mm,
nn::transformer builders, random::{randint,categorical,truncated_normal})
- fuse indexing presence-prove + position scans: index_op_to_array and
expand_ellipsis_operations halve the scan count on the hot
KV-cache slicing / attention-mask construction path
- drain VectorArray::try_into_array into stack [T; N] via MaybeUninit;
no Vec heap alloc on the quantized KV cache write path
- scalar-output compile/grad without Option unwrap: VectorArray::try_into_one
reads the lone Array from the C-vector directly; CompiledState gains
*_with_one paths used by 3 single-output CallMut impls
refactor:
- Entry API for get_mut_or_insert_with (optimizers update_single)
- propagate NumCast failures in arange/linspace as Exception::custom
instead of silent panic on NaN / out-of-range
- VectorArray::try_into_array fixed-N drain helper
- drop over-defensive Option unwraps across crate (no behaviour change)
- drop Option unwraps in random/adafactor
- drop unnecessary Result wraps; allow on trait-bound sites
- document remaining trait-impl unwraps via .expect("…") so panics are
self-documenting (PartialEq, Sum, Neg, IndexOp, IndexMutOp, Array::t)
feat:
- ops::contiguous wrapper (mx.contiguous) for callers that need an
explicit row-major copy
- new loader.rs with apply_post_load_memory_policy(): clear_cache then set_cache_limit to the resolved cap; called after eval_params in load_llama_model and load_qwen3_model - DEFAULT_CACHE_LIMIT_BYTES = 20 MB (matches mlx-swift LLM guidance) - set_cache_limit_override(n): programmatic override, first-call-wins - MLX_LM_CACHE_LIMIT_BYTES env var: runtime override - precedence: override > env > 20 MB default Loading n GB of safetensors stages through scratch buffers the MLX allocator parks in the reuse pool. Multi-model bench runs accumulate this dead memory. Decode-time reuse is small (<100 MB) and capping has no measurable perf cost on the upstream-supported models.
Trait: - fn attention(queries, keys, values, scale, mask) -> Array entry on KeyValueCache. Default: update_and_fetch + scaled_dot_product_attention. - resolve_sdpa_mask: None + n_q>1 is Causal; None + decode is no mask. - assert_mask_matches_keys (debug-only): catches [L, L] causal masks built without cache.offset(). - llama and qwen3 Attention::forward route through cache.attention(). Cache: - KVCache: pre-allocated [B, H, capacity, D] buffers. First call allocates init_capacity (default 64) tokens; overflows double the buffer (Vec-style geometric growth). update_and_fetch slice-writes the new tokens into the next [offset..offset+s] rows and returns graph-view Array slices over the populated [..offset] range. - Eliminates per-step concatenate_axis (O(N) per step → O(1) amortised). - Replaces ConcatKeyValueCache as the only KV cache type. Trait keeps is_quantized/group_size/bits as default no-ops for a future quantised-cache type to populate. Decode pipelining (added to the shared models::decode_step): - decode_step now async_evals its result: submit step N+1's forward + sample before the caller syncs N, so the GPU runs N+1 during N's sync. Generate::step delegates to it; the bench inherits the same overlap, so production and measurement still can't drift. - Generate gains DecodeFirst (primes the pipeline on first decode) + Decode states; prefill stays a single prompt forward. - inv_temp (1.0/temp) cached once on the iterator and threaded into decode_step; greedy (temp 0) takes the argmax path. - generate/mod.rs (disabled module): pipeline GenerateToken, read tokens back in one batched eval instead of try_item per step.
quantized_scaled_dot_product_attention had two dtype hazards (mlx is strict; a stray f32 promotes the whole graph and Metal rejects f64): - `queries * scale` (f32 scalar) promoted bf16/fp16 queries to f32, poisoning gemm + softmax for the rest of the forward. Stage scale into the query dtype: Array::from_f32(scale).as_dtype(q_dtype). - the bool-mask nucleus sentinel was Array::from_f64(finfo_min); f64 lands on the Metal stream and the dispatch fails. Build as f32, cast to the scores dtype. Both latent on this branch (only the quantized-cache SDPA branch hits them, and KVCache reports is_quantized() == false; the non-quantized path delegates scale to fast::scaled_dot_product_attention). Real bugs for any quantised KV cache that lands later.
The `generate/` module was disabled (`// pub mod generate`) and unused; its only consumers were the `ModelInput`/`ModelOutput`/`ModelInputBuilder` glue traits in lib.rs, which nothing else references. Remove both. The per-model `models::*::ModelInput` structs (the `Module::forward` input) are unaffected — distinct from the deleted `ModelInput` trait.
Parse config.json once into a typed ModelConfig with a Family enum
(serde tag = model_type). Unknown model_type fails with
`unknown variant` at parse instead of falling through a loader.
- config.rs: ModelConfig::from_dir + Family + quantization() that
reconciles `quantization` vs legacy `quantization_config` once.
- family.rs: EosSpec (int | [int] normaliser); envelopes carry
eos_token_id.
- models/{llama,qwen3}: ModelArgs loses model_type + quantization*
(now on the Family tag / outer struct), gains eos_token_id.
Loaders read off ModelConfig::from_dir; drop get_*_model_args +
the per-loader resolve_quantization duplication. Dead model_type()
accessor removed.
- error.rs: typed Error enum (Config/Shape builders + From<Error>
for Exception bridge).
The production runtime and the family-adapter dispatch land together, so
the model-local decode loop they replace never has to be introduced and
then removed across commits.
- LanguageModel/UserInputProcessor traits, ModelContext, generate(), the
chat-template + user-input + sampler surface (model_context.rs,
language_model.rs, lm_input.rs, user_input.rs, chat_template.rs,
sampler.rs).
- load() dispatch via the Family enum + per-family adapters
(llama/adapter.rs, qwen3/adapter.rs, family.rs).
- shared nn/ layer: AttentionInput + SwiGLU MLP hoisted out of the
byte-identical llama/qwen3 copies (nn/{attention_input,swiglu_mlp}.rs).
- decode_step is a ModelContext method over &mut dyn LanguageModel;
generate() and the bench both call it, so pipelining can't drift
between production and measurement. benches/lm_decode.rs runs one
family-agnostic path over ModelContext.
Breaking: mlx_lm::models::{llama,qwen3}::Generate and
models::{decode_step,sample_logits,inv_temp} removed; public entry is
mlx_lm::{load, generate}.
Restructure llama + qwen3 into one per-family layout so qwen3.5 (and future families) drop in identically. Pure relocation + dispatch convention — model graphs and internals unchanged, decode path untouched, bench-neutral. - models/<family>.rs -> <family>/text/model.rs; <family>/adapter.rs -> <family>/text/adapter.rs; ModelArgs split into <family>/text/config.rs. - Family: name() + as_llama()/as_qwen3(). Adapters take (cfg, args, dir) and stop re-reading config; load_context(cfg, dir) dispatches via as_<family>(). model_context::dispatch_load passes cfg through. - Drop the models/ module (code ported, not removed).
Groundwork for the qwen3.5 port: shared runtime + cache reshape, no
family wired in yet. KVCache + attention() carried verbatim (perf wins
preserved); FullAttnCache is a Standard-only KVCache passthrough.
- cache.rs -> cache/{mod,trait_def,kvcache,full_attn,options}.rs.
KVCache + KeyValueCache::attention() byte-identical (re-export keeps
crate::cache::KVCache path). FullAttnCache: Standard(KVCache) only,
zero per-step cost; Quantized arm + Pi rotation deferred to quant-KV.
CacheOptions/CacheKind + prefill-chunk helpers. build_rotation->None.
- language_model.rs: defaulted trait hooks prefill_chunk(_size),
has_mtp, try_mtp_decode, set_mtp_depth, set_cache_options. llama/qwen3
override nothing.
- model_context.rs: GenerateParams.disable_mtp, run_mtp_loop, use_mtp
branch, chunked run_prefill. decode_step path unchanged.
- .gitignore: un-ignore mlx-lm/src/cache/ (bare cache/ swallowed it).
QuantizedKVCache: affine-quant K/V with independent k_bits/v_bits, six
packed buffers grown geometrically. K is the softmax-sensitive tensor
(kept high), V tolerates low bits.
CacheKind::Quantized is { group_size, k_bits, v_bits } with presets
quantized_q8/k8_v8/k8_v4/q4 (DEFAULT_KV_GROUP_SIZE=64, MIN_K_BITS=8).
with_config clamps unsafe configs (k_bits up to MIN_K_BITS, v_bits down
to k_bits) with a stderr warning; unsupported bit-widths are a hard error.
Read dispatches on k_bits == v_bits: equal -> packed
quantized_scaled_dot_product_attention (shared bits); unequal -> dequant
each at its own bits then dense SDPA (one quantized_matmul can't mix
bit-widths). No rotation (TurboQuant lands later).
- utils::quantized_scaled_dot_product_attention queries Array -> &Array
(no per-call clone; only caller is the dead routing wrapper)
Add Compile::compile_with_id + allocate_compile_id; key the mlx-c compile cache on a process-wide monotonic id instead of TypeId. Two distinct fn-pointers cast to the same concrete signature share one TypeId, so keying mlx_detail_compile on the type id made the second compile() silently reuse the first function's compiled graph — e.g. an attention_gate of the same (&Array,&Array) signature returning sigmoid(output)*gate after a swiglu warmed the slot. Monotonic ids give one compiled-graph slot per call regardless of source type; a stable per-op id (allocate_compile_id) lets many module instances share one compiled Metal kernel.
Interactive REPL against any mlx_lm checkpoint (ported from realign,
server path dropped). Streams assistant output through a small
think/answer colouriser: reasoning (<think>…</think>, tags stripped)
renders dim, the answer bold green, and the readline prompt bold cyan
so user input reads distinctly. Resets the KV cache per turn and
re-renders full chat history each request.
Flags: --model, --temperature/--top-p, --max-tokens, --think
(on|off|default), --prefill-chunk-size, and the KV-cache controls:
--kv-cache standard|q8|q4 preset (default standard)
--k-bits / --v-bits per-tensor bit override (forces quantised;
unset tensor follows the preset, else the
other bit; bare --k-bits 8 -> k8/v8)
--kv-group quantisation group size (default 64)
The resolved K/V/group config is logged to stderr before generation.
- new workspace member examples/chat (publish = false)
- no markdown/table rendering; deps trimmed to argh/rustyline/anyhow/
serde_json/env_logger
Dense text path for the Qwen3.5 hybrid family: full-attention layers + gated-delta-net (GDN) linear-attention layers, dispatched per config.json layer_types. Loads + decodes a real dense checkpoint (text-only load from a VLM checkpoint ignores vision-tower weights). - qwen3_5/text/: config (text_config envelope), mrope, attention + SwiGLU MLP, GDN block (depthwise Conv1d + cached Metal recurrent scan), hybrid LayerCache, weights, dense adapter. - config: Family::Qwen35 + as_qwen35()/name(); lib + dispatch arm. - activations: plain swiglu/attention_gate. - loader: load_tokenizer/list_shards/rewrite_quantised_keys. - mlx-tests: dense GDN-hybrid load+decode e2e (ignored; needs MODEL).
- Qwen35MoeBlock: shared+routed DeepSeek-style MoE over the qwen3.5 hybrid GDN + full-attention spine; sigmoid-gated shared expert. - nn::switch: SwitchLinear/QuantizedSwitchLinear (gather_mm/gather_qmm), SplitSwitchFfn + SwigluActivation; expert-id sort above threshold. - quantization: per-tensor overrides + for_path; QuantMode enum. mlp.gate/shared_expert_gate requantised to 8-bit when overridden. - MTP self-speculative decode: adapter_moe rejection loop + sampling top_p_keep_mask; SamplerState gains sampler()/masked_log_probs(). - Family::Qwen35Moe variant (model_type qwen3_5_moe) + dispatch.
- qwen3_5 dense GDN-hybrid (Qwen3.5-4B q8/q4): GDN scan kernel + full-attn hybrid spine on the production decode path. - qwen3_5_moe (Qwen3.6-35B-A3B q8): gather_qmm experts + MoE routing. - maybe_bench_mtp A/B: generate() MTP-on vs -off throughput at temp>0 (top-p), keyed on emitted tokens. - Full set only (heavy models); each self-skips when checkpoint absent.
Add the cross-cutting public surface a vision adapter consumes, gated on
a new default-on `image` cargo feature (text-only builds compile with
--no-default-features, no codec dep):
- LMInput.image: Option<ProcessedImage> + ProcessedImage {pixels, grids}
- UserInput.images: Vec<Image> + with_images builder; Image enum
(Decoded(DynamicImage) | Pixels {array, grid})
- Error::OutOfBounds + Error::out_of_bounds (multimodal index checks)
- lib re-exports: ProcessedImage, Image
Runtime-inert: no VLM code yet. Text adapters set image: None; the whole
vision surface compiles out feature-off.
Add the qwen3.5 VLM under qwen3_5/image/ (feature-gated on `image`): - vision.rs: Qwen3-VL ViT tower (Conv3d PatchEmbed, 2D-rope cu_seqlens attention, bilinear pos-embed interp, PatchMerger). - processor.rs: Qwen35ImageProcessor (smart_resize, patchify) over the image crate. - multimodal.rs: merge_input_ids_with_image_features + mrope get_rope_index_single_batch/_batched + pack_position_ids. - weights.rs: load_full_model splits LM + vision-tower params. - adapter.rs: Qwen35VlmAdapter wraps Qwen35DenseAdapter; Qwen35Processor. Re-adds the seams stripped in the dense/MoE commits, now at their VLM consumer (all #[cfg(feature = "image")]): Bucketed::Vision routing in text::weights (text loaders drop it), Qwen35Model/Decoder forward_embeds + the inputs_embeds path, Qwen35DenseAdapter cursor/rope_delta + prefill_embeds. Family::Qwen35Vl + dispatch; load_context probes preprocessor_config.json to route VL checkpoints, text-only fallback when the feature is off. mlx-tests gains a feature-gated, ignored VL e2e (needs a dense qwen3.5-VL checkpoint via MODEL + IMAGE).
2× ring-buffer KV cache for sliding-window attention (Gemma 3/4 sliding layers). Layout [B,H,keep+2*window,D]: keep prefix is write-once, the rotating region wraps with O(window) amortised-O(1)/token compaction so steady-state decode is one try_index_mut write + one contiguous view. Multi-token prefill returns old_window ++ new so each new token attends the full window without a mask recompute. Implements KeyValueCache (offset/max_size/update_and_fetch; attention via the trait default). inherent trim/is_trimmable. No steel-prefill kernel (uses fast::sdpa) and no state serialization yet — added with a consumer.
Hybrid sliding/global attention (per-layer head_dim, q/k/v norms), proportional partial-rotation rope, GeGLU MLP, four norms per layer, embedding scaling, final-logit soft-capping, tied embeddings. Validated against gemma-4-31b (dense, 60 layers, sliding_window 1024). Wired into Family::Gemma4 + load dispatch. Shared additions (reusable beyond gemma): - nn::RmsNormNoScale (rms_norm with no learnable gain) - activations: geglu / logit_softcap / residual_add_scale compiled caches - cache::effective_prefill_chunk_opt (windowed prefill cap) MoE, per-layer-input embeddings (E2B/E4B), KV-sharing, and vision are deferred; each re-adds at its own consumer. The sliding cache's single forward must not exceed sliding_window — the adapter's prefill_chunk_size enforces it via effective_prefill_chunk_opt.
The fast tokenizer's post_processor omits BOS for several mlx-community
conversions (Gemma drops it entirely; the single-template adds no
special tokens), so raw-text prompts reached the model without <bos>.
BOS-sensitive families collapse without it — Gemma 4 emitted degenerate
repetition (" France is France is...") instead of coherent text.
- resolve_bos_id: read add_bos_token + bos_token from tokenizer_config,
resolve the id via the tokenizer; honour explicit add_bos_token:false,
else prepend when a bos_token exists (HF default for Gemma/Llama 3).
- TextOnlyProcessor gains bos_id; prepend in prepare unless already
present (chat templates emitting {{ bos_token }} keep one BOS).
- Wire through all five text adapters. Qwen (bos_token:null) -> None,
byte-identical output, no regression.
Dual-branch decoder layer for the enable_moe_block variant: the dense
GeGLU MLP and a 128-expert top-8 routed FFN run in parallel and sum.
Reuses the shared SplitSwitchFfn + router_topk kernel; the experts bind
the checkpoint's split experts.switch_glu.{gate,up,down}_proj keys.
- nn/switch.rs: GegluActivation (SwitchActivation over the geglu cache).
- gemma4/text/moe.rs: Router (rms_norm with scale*hidden^-0.5 -> proj ->
router_topk -> *per_expert_scale[idx]) + Experts
(SplitSwitchFfn<GegluActivation>).
- DecoderLayer: enable_moe gate with Option router/experts + 3 extra
norms (post_ff_1, pre_ff_2, post_ff_2). forward_layer now returns
AttentionOut and takes shared_kv/offset/per_layer_input (None in the
dense+MoE base) — the one forward-compat hook for the E2B/E4B follow-on.
- Single Family::Gemma4 dispatch: dense vs MoE decided per-layer in
DecoderLayer::new from enable_moe_block. No second adapter; the loader
binds split experts/router via the existing quantised-key rewrite.
- bench: gemma4 26b_a4b q8/q4 cells. e2e: gemma4_moe_e2e (Paris).
The E2B/E4B variants add two mechanisms over the dense/MoE base, both config-gated (inert when their fields are 0, so 31b/26b are unchanged): - Per-layer-input embeddings (hidden_size_per_layer_input > 0): a separate embedding plus a projection of the main hidden, scaled/normed/averaged into a [B,L,num_layers,pl] tensor sliced per layer; each layer gates it in via gelu(gate(h)) * pl_in -> proj -> norm -> add (between the FFN residual and the layer-scalar multiply). Scales staged in h dtype. - KV sharing (num_kv_shared_layers > 0): the last N layers own no K/V projection and reuse a prior same-kind layer's (k, v) + offset. Attention K/V/norm become Option (has_kv gate); compute_previous_kvs builds the source-index table; the model threads each layer's fetched KV through intermediates into its downstream sharer; make_caches yields None slots for shared layers; the loader drops their self_attn.k_*/v_* keys. - Mlp::new takes the effective intermediate width — doubled on KV-shared layers when use_double_wide_mlp (E2B). Reuses the C-A forward_layer signature (shared_kv/offset/per_layer_input) unchanged. Validated e2e (chat -> Paris) on gemma-4-e4b-it-8bit (42 layers, 18 shared) and gemma-4-e2b-it-8bit (35 layers, 20 shared, double-wide).
- gemma4/image/vision.rs: PatchEmbedder (Linear patchify + per-axis position table), bidirectional VisionAttention (q/k norm, param-free v-norm, 2-D RoPE, scale=1), GeGLU VisionMlp, VisionEncoderLayer (gemma sandwich norms), VisionPooler (k×k avg-pool ×√hidden), VisionModel (standardize), EmbedVision (RMS-noscale → quantized Linear) - gemma4/image/config.rs: VisionConfig (gemma4_vision) - gemma4/image/multimodal.rs: stitch over qwen merge helper - text.rs: forward_embeds (Model + Gemma4TextModel) + embed_tokens accessor; extract embed_scaled/forward_from_hidden/apply_head; text embeds carry embed_scale, stitched vision features do not - text/weights.rs: expose rewrite_outer_key + is_shared_kv_layer_key Vision tower weights bf16; projector quantized. Synthetic shape tests.
- text/config.rs: vision_config (image-gated) + image/boi/eoi token ids on gemma4 ModelConfig - gemma4/mod.rs: route vision_config + processor_config.json checkpoints to the VLM adapter (image feature); else text-only - image/processor.rs: Gemma4ImageProcessor — aspect-preserving resize to a pooling·patch-divisible grid, rescale to [0,1], channels-first - image/weights.rs: load_full_model — bucket keys into text / bf16 vision tower / quantized projector; quantize text + embed_vision only - image/adapter.rs: Gemma4VlmAdapter (tower → projector → embed_scaled → stitch → forward_embeds; image-token ids masked to 0 for per-layer inputs) + Gemma4Processor (chat render, <|image|>→boi+image×N+eoi, token-count assert) - vision.rs: standardize after the pooler's √hidden scaling (per-token bias, not per-patch) — the order that matches the reference - gemma4_vision_e2e.rs: ignored MODEL+IMAGE caption smoke test Validated on gemma-4-31b-it-4bit: caption matches mlx-vlm reference.
- --image <path>: decode + send as one user_with_image turn, then exit (skips the REPL). VLM dispatch is automatic once UserInput carries the image; --prompt is the instruction (default blank). - factor think-kwarg + sampling/params resolution into apply_think / build_params, shared by the one-shot path and the REPL loop. - image dep (png/jpeg) for image::open. Validated on gemma-4-31b-it-4bit: --image + prompt -> correct caption; blank prompt -> model captions unprompted; text REPL unchanged.
Opt-in `audio` feature (implies `image`). Fixes e2b/e4b load: the audio tower keys are now bound (audio on) or dropped (audio off) instead of landing unbound in the text bucket. - gemma4/audio/config.rs: AudioConfig; ModelConfig gains audio_config + audio/boa/eoa token ids. - gemma4/audio/clippable.rs: ClippableLinear (Linear + live input/output clamps; audio sets use_clipped_linears=true). - gemma4/audio/encoder.rs: SubSampleConvProjection (2x Conv2d subsample), ConformerFeedForward (macaron), AudioAttention (chunked local attn, relative-position bias + relshift, per-dim-scale softplus, logit softcap, f32), ConformerLightConv1d (depthwise causal-conv GLU), ConformerBlock, AudioEncoder, EmbedAudio (RMS-noscale -> quantized Linear 1536->text_hidden). - gemma4/audio/multimodal.rs: stitch over the shared masked-scatter. - vlm/weights.rs: generalized VLM loader returns LoadedTowers with an optional audio tower+projector; audio bf16, projector quantized; vision clip buffers dropped, audio clip buffers kept (live). - vlm/adapter.rs: adapter holds the optional audio tower; prepare() stitches image and/or audio; image+audio token ids masked to 0 for PLE. - LMInput gains an audio carrier (ProcessedAudio log-mel). Synthetic shape/clip unit tests. Validated: e4b loads clean with audio on (tower binds) and off (keys dropped). Front-end + input wiring follow.
- gemma4/audio/feature.rs: log_mel front-end (16kHz mono -> [1,T,128]) -- framed periodic-Hann + rfft + HTK mel filterbank + log; num_audio_tokens derives the soft-token count from the conv-subsample frame math. - vlm/adapter.rs: Gemma4Processor audio branch (log_mel, <|audio|> -> boa+audio*N+eoa expand, token-count assert) -> ProcessedAudio carrier; modality-generic marker expansion shared with images. - chat_template.rs: ContentPart::Audio + ChatMessage::user_with_audio. - user_input.rs: Audio carrier + with_audio; lm_input ProcessedAudio. - examples/chat: --audio one-shot (16kHz mono WAV via hound; behind the chat `audio` feature). Tests (committed sine fixture + golden .f32, CI-safe, no model): - gemma4_audio_logmel: log_mel vs golden (front-end numerics lock). - gemma4_audio_encoder: seeded synthetic-weights encoder vs golden (encoder math lock). - vlm/weights.rs audio_encoder_real_weights_stats (#[ignore], MODEL): real-weight load-path lock (clip buffers, key binding). - gemma4_audio_e2e (#[ignore], MODEL+AUDIO): transcription smoke test. Validated on gemma-4-e4b-it-8bit: transcript matches the reference; log_mel <0.003 max-abs on the committed fixture.
Author
|
I went well overboard with this. Most of it is reviewed multiple times by myself and a few AI agents. The last 4-5 commits could use a finer eyeball I think. The series has been project for the last few weeks since I needed some particular features for my own goals. I've done my best to maintain benchmarks, and have quite thoroughly tested the bulk of all the series and note some steady perf increases in most things except for qwen35/gemma4 as these are a reimplemented of work I did in another branch that was very throughly tweaked. |
Extract the family-agnostic speculative-decode primitives so a second family (gemma4 drafter) can reuse them; zero behavior change. - crate::speculative: top_p_keep_mask + sample_draft/accept_mask/ resample_on_reject (was qwen3_5/text/sampling.rs + adapter_moe.rs) and a generic CacheSnapshot<T> (was the qwen-local two-cache struct). - qwen3_5 MoE adapter calls the shared helpers; the two-cache snapshot becomes two single-cache guards (committed/rolled together — same behavior). - KeyValueCache::current_kv: non-mutating dense (k,v) over the cached history (None when empty; quantized dequants). KVCache/RotatingKVCache/ FullAttnCache/QuantizedKVCache + gemma LayerCache. RotatingKVCache promotes snapshot_window. Gate: qwen MoE mtp_greedy_matches_plain_decode unchanged; current_kv unit test vs update_and_fetch.
Load a separate `gemma4_assistant` drafter checkpoint alongside a gemma4 target and run speculative decode: the drafter (a 4-layer Q-only gemma4 stack) borrows the target's per-type cached K/V + last hidden, drafts γ tokens, and the target verifies them in one parallel forward. Greedy is byte-exact vs plain decode (the spec-decode guarantee). - gemma4/mtp/config.rs: DrafterConfig (`gemma4_assistant`; nested gemma4 TextConfig, backbone width, centroid flags, per-size depth). - gemma4/mtp/drafter.rs: Drafter = gemma4::text DecoderLayers (all KV-shared → Q-only) + pre_projection(2*backbone→draft) + post_projection(draft→backbone) + own tied embed/head. Borrows target K/V via the existing `shared_kv` seam (cache None — read-only). - gemma4/mtp/centroid.rs: MaskedEmbedder sparse lm head (E2B/E4B `use_ordered_embeddings`): centroid scores → top-k clusters → exact logits for those tokens via `token_ordering`, scatter to vocab. - gemma4/mtp/decode.rs: snapshot → draft γ (read-only) → one verify forward → accept prefix + bonus (greedy) or rejection-sample, rolling back + re-committing the accepted prefix on a partial reject. - gemma4/mtp/weights.rs: assistant loader; hydrates the token_ordering buffer (not a #[param]). - text.rs: forward_hidden_and_logits + embed_scaled_token (the drafter concat uses the embed_scale×√hidden token embedding the assistant was trained on). config.rs: deserialize_softcap accepts null (drafter sets final_logit_softcapping=null). - adapter.rs: Gemma4Adapter holds Option<Drafter>; has_mtp/try_mtp_decode/ set_mtp_depth; captures prev_hidden in prepare/step/prefill_chunk. - model_context.rs: load_with_drafter(dir, draft_dir); gemma4-only. - gemma4_mtp_e2e.rs: greedy parity (drafter == plain greedy) — the correctness gate. Validated on gemma-4-e4b-it-8bit (target) + gemma-4-E4B-it-assistant: greedy with/without the drafter match token-for-token.
- --draft-model <dir>: load a gemma4 MTP assistant (drafter) alongside the target via load_with_drafter; speculative decode runs automatically. - --mtp-depth <n>: override drafter depth γ (default per assistant size). Validated: gemma-4-e4b-it-8bit + gemma-4-E4B-it-assistant one-shot generates correctly with the drafter loaded.
- DrafterConfig gains optional quantization (from assistant config.json) - load_drafter applies try_into_quantized before binding, matching the rewrite_quantised_keys inner.weight key path - prior loader assumed bf16 drafters; quantized assistant checkpoints (e.g. *-assistant-8bit) failed to bind embed_tokens/proj keys
- new Family::Gemma4Unified variant + dispatch (model_type gemma4_unified)
- gemma4_unified/{config,adapter,mod}: dense text + MTP speculative decode
- reuses gemma4::text::{Model,cache} + gemma4::mtp wholesale; the unified
text backbone + assistant drafter are structurally identical
- gemma4 text loader: load_model takes &TextConfig (shared by both
families); drop vision_embedder keys (encoder-free MM front-end)
- gemma4 text loader honours per-tensor quant overrides (mlx-community
4bit ships MLP gate/down/up_proj at 8bit over a 4bit body)
- extract requantise_linear to crate::quantization; qwen3_5 MoE reuses it
- encoder-free vision/audio front-ends deferred to follow-on milestones
- image/{config,embedder,processor,weights,adapter}: encoder-free vision
path. No SigLIP tower — patch_ln1 → patch_dense → patch_ln2 →
factorised-2D-posemb → pos_norm → RMSNorm(no-scale) → projection
- processor extracts model_patch_size (48px) HWC pixel blocks + 2D
position ids, pads to num_soft_tokens with -1; equals the HF
convert_image_to_patches + patches_merge result (cross-checked in test)
- VLM adapter strips padding patches, stitches soft tokens into
image_token_id slots, decodes via Model::forward_embeds
- MTP composes with vision: image affects prefill only, so the drafter
decodes off the populated KV cache (forward_embeds_hidden_and_logits
seeds prev_hidden). --draft-model + --image now both active
- mod.rs routes to the VLM adapter when vision_config + processor_config
present (drafter optional)
- ProcessedImage gains optional position_ids (encoder-free vision needs
per-patch 2D ids; tower families pass None)
- rewrite_quantised_keys also redirects <prefix>.bias → inner.bias for
quantised linears with a bias (patch_dense); was weight-only
- audio/{config,feature,embedder,weights}: encoder-free audio path. No USM
Conformer, no mel — raw 16kHz waveform framed into audio_samples_per_token
(640 = 40ms) rows, projected via RMSNorm(no-scale) → linear into text space
- feature: pad + reshape waveform to [T, 640]; 25 tokens/sec
- VLM adapter handles audio alongside image: encode_audio + stitch into
audio_token_id slots; composes with MTP (audio is prefill-only)
- processor frames clips, expands <|audio|> markers to boa+audio×N+eoa,
asserts token count; render_prompt/count_match gain audio
- ProcessedAudio.mel renamed to .features (mel for gemma4 USM tower, raw
frames for unified); gemma4 vlm adapter updated to match
- chat --audio one-shot works for the unified family
- in-tree converter: load bf16 shards → per-family Rewriter → quantise → sharded safetensors + index + quantised config.json + tokenizer assets - Rewriter trait + QuantClass (Skip/Body/Pinned) drive per-tensor rules; unknown shapes hard-error rather than silently dropping tensors - runner evals per shard to cap peak at ~1x model size - qwen3_5 family wired (dense + MoE incl. MTP, gate-up split) - --verify round-trips the output through mlx_lm::load + a short generate - workspace deps: add serde/serde_json/anyhow/argh/env_logger
- Gemma4UnifiedRewriter: uniform body-quant (no per-tensor pins) → a true q4, unlike mlx-community's mixed 4bit (MLP pinned q8) - body-quant projection weights (attn q/k/v/o, MLP gate/down/up, embed_tokens, encoder-free patch_dense + embed_vision/audio projections); skip norms (RMSNorm + vision LayerNorm), layer_scalar, pos_embedding, biases - keys pass through unchanged (already the loader's expected form) - drop_mm flag (default false) → full-MM q4; true → text-only - --family gemma4_unified wired into the bin === downloading bf16 mirror, check progress === du -sh /tmp/full-models/gemma-4-12B-it-bf16 2>/dev/null; tail -1 /tmp/bf16_dl.log 2>/dev/null
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Infra / core
KV cache
Runtime / arch
qwen3.5
gemma4
examples/chat
Fixes: BOS prepend in TextOnlyProcessor.
Two new model families (qwen3.5, gemma4) full text+VLM, gemma4 audio, MTP speculative decode end-to-end.