Gemma 4 E2B QAT 4-bit assistant — experimental MLX patches
This is a self-contained copy of the
mlx-community/gemma-4-E2B-it-qat-assistant-4bit MTP drafter, plus three
experimental runtime patches for MLX-VLM 0.6.13:
qat_4bit_assistant_patch.py gathers tied embedding rows through MLX's
public QuantizedEmbedding path. Without it, the released 4-bit assistant
is treated as if its packed weights were a dense embedding matrix and
crashes during drafting.qat_structured_mtp_patch.py applies caller-provided JSON-schema masks to
target verification tokens during MTP. Draft proposals stay unconstrained,
but a proposal is accepted only when it equals the grammar-constrained
target choice.qat_gemma_target_verify_patch.py ports MLX-VLM's Qwen MTP parity design:
AR-equivalent quantized projection/attention kernels plus physical Gemma
rotating-cache ordering and rejected-slot rollback.
The weights themselves are unchanged. The first two patches are source-hash
gated. The block-verification prototype relies on private 0.6.13 APIs and must
be revalidated, not carried forward blindly, on any other MLX-VLM version.
Important status
Experimental; not production-ready. The structured bridge guarantees the
tested JSON grammar, but it does not repair a separate greedy-parity defect in
MLX-VLM's Gemma 4 MTP implementation.
On 30 sealed EXAMI extraction calls at temperature 0:
| runtime | valid schema | byte-identical to constrained AR | effective tok/s |
|---|
| constrained QAT AR | 30/30 | control | 38.0 |
| patched constrained MTP, block 4 | 30/30 | 13/30 | 73.1 |
| patched constrained MTP, block 2 | 28/30 | 11/30 | 51.5 |
Block 4 is about 1.9x faster, but speculative decoding at temperature 0 must
be byte-identical to the target's greedy output. This repository therefore
documents a working structured-MTP prototype, not an approved inference path.
See PATCH_RESULTS.json for the measured values.
The block-level target-verification patch now avoids full-ring snapshots. It
rewinds and replays only overwritten physical-ring slots for shared-KV layers.
On a fresh matched 30-call gate:
| runtime | byte-identical to ordinary AR | effective tok/s |
|---|
| ordinary QAT AR | control | 41.3 |
| exact 4/5/8-bit block verify + batched ring journal | 30/30 | 53.05 |
This is not serial target generation: the target still verifies a speculative
block in one model forward. Adding the singleton-exact 8-bit QMV kernel removes
the expensive four-singleton fallback required by this mixed-precision QAT
target. The verifier now snapshots every overwritten slot in one GPU
synchronization per layer. MTP was 28.3% faster than AR (267.84 vs 343.70 s).
Both runtimes reproduced the same one 1,400-token truncation (29/30 valid JSON),
so that format failure is not introduced by MTP.
A subsequent sealed, card-level blind pass found no reliable semantic-quality
difference from constrained AR:
| runtime | cards | unique TRUE-defective cards | exact clean cards | TRUE incidence | compound |
|---|
| constrained QAT AR | 271 | 12 (4.43%) | 95.57% | 4.80% | 7.38% |
| patched MTP block 4 | 263 | 13 (4.94%) | 95.06% | 5.70% | 7.60% |
The unique-defect delta was +0.51 percentage points for MTP, with a paired
source-chunk bootstrap 95% interval of −0.97 to +2.25 pp. This was one
within-run-blind Codex reader, not independent two-reader confirmation. The
audit concerns the earlier divergent fast path, where only 13/30 outputs
matched AR. The exact verifier below is byte-identical to AR by construction.
Install
Use a fresh Python 3.11 environment on Apple Silicon:
1python3.11 -m venv .venv
2.venv/bin/pip install -r requirements.txt
The example uses the compatible target separately; by default it downloads
mlx-community/gemma-4-E2B-it-qat-4bit. Access remains subject to the upstream
Gemma license and model terms.
Run the structured-MTP prototype
1.venv/bin/python example_structured_mtp.py \
2 --prompt 'Return a JSON object whose answer explains why the sky appears blue.'
Pass a local or alternative compatible QAT target with --target.
Run the block-level parity prototype separately:
1.venv/bin/python example_block_parity_prototype.py \
2 --prompt 'Return a JSON object whose answer is Paris.'
For an image request, this repository includes the required torch-free Gemma
visual-processor registration and a separate example:
1.venv/bin/python example_image_structured_mtp.py \
2 --image /path/to/image.png
The image path was verified on the EXAMI logo: constrained AR and patched MTP
returned the same 40-token description byte-for-byte, identifying the stylized
E, graduation cap, pencil, reddish-brown foreground, and black background. This
is a functional multimodal smoke, not a powered image-quality benchmark.
Files
model.safetensors and the tokenizer/config files: original 4-bit QAT MTP
assistant artifact.qat_4bit_assistant_patch.py: quantized tied-embedding compatibility bridge.qat_structured_mtp_patch.py: sequential target-side grammar masking during
speculative verification.qat_gemma_target_verify_patch.py: exact 4/5/8-bit block verifier plus
batched physical-ring rollback; 30/30 parity, 53.05 effective tok/s.example_structured_mtp.py: minimal end-to-end example.example_block_parity_prototype.py: runnable correctness prototype.example_image_structured_mtp.py: end-to-end single-image example.qwen_vlm_torchfree_patch.py: registers MLX-VLM's NumPy/PIL Gemma visual
processor without PyTorch; its filename is historical and it supports Gemma.PATCH_RESULTS.json: measured correctness and speed gates.requirements.txt: exact tested MLX stack.
Uploading to Hugging Face
The included .gitattributes already routes Safetensors and tokenizer data
through Git LFS. Upload this directory as the root of a model repository. Do
not remove the experimental warning or describe the fast MTP path as lossless.
The exact verifier passed 30/30 byte identity and measured 28.3% faster than
AR. Keep the experimental label until independently replicated/upstreamed.
Upstream context
- MLX-VLM 0.6.13 currently rejects structured output together with speculative
decoding in its server.
- vLLM supports Gemma 4 assistant checkpoints through its MTP path and permits
structured-output requests, but it is a CUDA/PyTorch serving stack rather
than an MLX/Apple-Silicon replacement.
- TensorRT-LLM also implements guided plus speculative decoding and documents
grammar advance/rollback as part of target verification.
- MLX-VLM's Qwen MTP parity repairs (#1188/#1210) use dedicated target-verify
kernels; Gemma does not yet have an equivalent upstream implementation.
This repository follows the upstream google/gemma-4-E2B-it model card and
Gemma license terms.
