Bei0001/gemma-4-E2B-it-qat-assistant-4bit-mlx-patched

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:

1. 1.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.
2. 2.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.
3. 3.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:

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:

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:

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:

python3.11 -m venv .venv
.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

.venv/bin/python example_structured_mtp.py \
  --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:

.venv/bin/python example_block_parity_prototype.py \
  --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:

.venv/bin/python example_image_structured_mtp.py \
  --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.