VoxtralRealtime: load checkpoints with a quantized tied embedding

[T0mSIlver]

Allows loading Voxtral Realtime models whose tied embedding is quantized. Today such models fail to load — which leaves the most expensive per-token operation of the decode loop stuck in fp16 (see the numbers below).

Problem

Voxtral Realtime has no separate LM-head weight: the decoder reuses the token-embedding matrix (tok_embeddings) both to embed input tokens and to project hidden states to logits ("tied embedding").

A converted checkpoint can store this matrix in one of two forms:

• fp16: a single tok_embeddings.weight tensor. This is what the current mlx-community conversions ship, and it loads fine.
• quantized: three tensors — a bit-packed weight plus scales and biases. This is the form MLX's standard quantization path (nn.quantize, which includes embeddings) produces, so converted checkpoints commonly ship it for other model families.

The quantized form cannot be loaded. During model construction the loader runs a quantize pass that turns each Linear/Embedding into its quantized counterpart when the checkpoint has scales for it — but shouldQuantize hard-excludes tok_embeddings, so the module always stays a plain Embedding that expects exactly one fp16 tensor. Weight loading with verify: .all then rejects the checkpoint's tok_embeddings.scales/.biases with an UpdateError — a plain Embedding has no such parameters.

Why it matters

For the 4B model the tied head is a 131072×3072 matrix: 768 MiB in fp16, read in full for every decoded token's logits row. In live streaming on an M1 Pro that fp16 gemv is the single largest phase of the decode loop. With a 4-bit/group-64 quantized tied embedding (measured, same machine, same audio):

• LM-head projection: ~30 ms → ~3 ms per decoded token (quantized-matmul fast path)
• end-to-end 100 ms-cadence streaming step: 92–104 ms → 64–74 ms
• resident weight memory: ≈ −530 MiB (768 MiB fp16 → ≈ 240 MiB packed weight + scales/biases)
• transcription quality held level on a 156-case English/French dictation eval with heavy technical vocabulary: 50/156 hard cases fully passing vs 49/156 for the fp16-embedding checkpoint, mean word accuracy 0.73 vs 0.72 on its stress subset

(The streaming numbers are from one realtime workload on one machine; the per-token LM-head saving applies to any decode loop over this model.)

Why the exclusion existed

The exclusion and the .scales gate both date from the original model port (#52), and it looks deliberate:

1. It matches what the reference conversions leave unquantized — norms, conv stem, adapter projection, embeddings.
2. More importantly, two decode paths read the raw weight tensor directly: embedToken gathers a row with weight[tokenId], and logits computes matmul(h, weight.T). On a QuantizedEmbedding the weight property is the bit-packed tensor, so if a quantized embedding had ever been structured, both paths would have consumed packed bits and produced garbage tokens with no error. Refusing to load was the safer failure mode.

This PR fixes the two paths, at which point the exclusion only blocks legitimate checkpoints.

Change

Three pieces. For fp16-embedding checkpoints the gather is bit-identical and the projection is the same contraction (within kernel rounding — the tests pin both):

1. embedToken gathers through the module call instead of indexing weight — on a QuantizedEmbedding that dequantizes the row, on a plain Embedding it is the same gather.
2. logits projects through Embedding.asLinear — the same contraction for a plain embedding, the quantized-matmul path for a quantized one.
3. shouldQuantize no longer excludes tok_embeddings. The quantize pass is still gated on the checkpoint shipping <path>.scales, so fp16-embedding checkpoints are structured exactly as before; nothing is quantized that the checkpoint didn't quantize.

Tests

Tests/VoxtralRealtimeTiedEmbeddingTests.swift:

• Plain embedding: module-routed embedToken/logits match the raw-weight formulas they replaced (gather exactly; projection within kernel rounding).
• Quantized embedding (via the standard quantize(model:filter:) machinery): embedToken returns exactly the dequantized row, and the tied projection matches the dequantized-weight contraction within the quantized kernel's accumulation error.

Also validated end-to-end on device: a Voxtral-Mini-4B checkpoint with a 4-bit/g64-quantized tied embedding loads and transcribes real audio correctly through the streaming session (the eval above).

[T0mSIlver]

For anyone wanting to try this: I've published a ready-made checkpoint with the 4-bit/g64-quantized tied embedding — https://huggingface.co/T0mSIlver/Voxtral-Mini-4B-Realtime-2602-4bit-qhead (otherwise the same conversion as mlx-community's 4-bit).

Worth underlining that this is a big realtime win for quantized-model users, not just a compatibility fix: the tied head is the decode loop's single largest phase, and quantizing it takes the LM-head projection from ~30 ms to ~3 ms per decoded token on an M1 Pro (~30 ms off every 100 ms streaming step, ~530 MiB less resident memory, transcription quality level on our dictation eval). And since MLX's standard nn.quantize path already produces quantized embeddings for other model families, any Voxtral checkpoint converted that way just works once this loads.