Speculative Decoding Failure with Gemma-4-31B and Assistant Drafter

[nunez]

Summary

Speculative decoding fails with an MLX error: expected a non-empty mlx_array when using a Gemma-4-31B-Instruct model with a corresponding assistant drafter model. Disabling the drafter allows the server to start successfully.

Environment

• OS: macOS
• MLX Serve Version: (Check with mlx-serve --version)
• Main Model: mlx-community/gemma-4-31b-it-bf16
• Drafter Model: mlx-community/gemma-4-31B-it-assistant-bf16
• MLX Core App: /Applications/MLX Core.app

Command That Fails

"/Applications/MLX Core.app/Contents/MacOS/mlx-serve" \
  --model ~/.omlx/models/mlx-community/gemma-4-31b-it-bf16 \
  --serve --port 8080 --host 0.0.0.0 --log-level info \
  --model-dir ~/.omlx/models/mlx-community \
  --metrics \
  --no-vision \
  --drafter ~/.mlx-serve/models/mlx-community/gemma-4-31B-it-assistant-bf16 \
  --draft-block-size 4 \
  --max-concurrent 4 \
  --prefix-cache-entries 8 \
  --prefix-cache-disk 10GB \
  --temp 0.8 \
  --top-p 0.95

Error Output

Model: gemma4 (60 layers, 5376-dim, head_dim=256, 32h/16kv, 0-bit affine quant)        
Loading tokenizer...  
  parsed tokenizer (30 MB) in 105ms
  loaded 262144 vocab entries in 6ms
  loaded 514906 merges in 15ms
Tokenizer loaded: 262144 vocab, 514906 merges, 24 special tokens (SentencePiece BPE)
User turn marker: "<|turn>user\n" -> 3 tokens
[registry] max_resident_models=3, max_resident_mem=375.0 GB
[preflight] weights ~58.25 GB, available 487.12 GB
Loading model-00012-of-00012.safetensors...
... (all 12 shards loaded)
Loaded 832 weights from 12 file(s)
Precomputing layer weights...
Batch eval all weights: 2069ms
GELU compiled (kernel fusion enabled)
GeGLU compiled (kernel fusion enabled)
Softcap compiled (kernel fusion enabled)
Loading model.safetensors...
Loaded 48 weights from 1 file(s)
[drafter] loaded 4 layers, vocab=262144, hidden=1024 → backbone 5376
MLX error: expected a non-empty mlx_array at /private/tmp/mlx-c-20260424-8100-j29lpv/mlx-c-0.6.0/mlx/c/ops.cpp:3701

Analysis

The error occurs immediately after the drafter model is loaded:

[drafter] loaded 4 layers, vocab=262144, hidden=1024 → backbone 5376
MLX error: expected a non-empty mlx_array

Key observations:

1. Dimension Mismatch: The drafter has hidden=1024 while the backbone has hidden=5376. This is expected for speculative decoding (the drafter is smaller), but the error suggests the tensors being passed between them are incompatible.

2. Empty Array Error: The expected a non-empty mlx_array error at ops.cpp:3701 typically occurs when:

   • A tensor operation receives an uninitialized or empty tensor
   • A shape mismatch causes an operation to produce an empty result
   • An embedding lookup or matrix multiplication fails due to incompatible dimensions

3. Potential Causes:

   • The drafter model's embedding layer may not be properly initialized
   • The speculative decoding kernel may not handle the dimension jump from 1024 → 5376 correctly
   • The gemma-4-31B-it-assistant-bf16 model may be malformed or incompatible with the current MLX version

Steps to Reproduce

1. Download the main model: mlx-community/gemma-4-31b-it-bf16
2. Download the drafter model: mlx-community/gemma-4-31B-it-assistant-bf16
3. Run mlx-serve with the --drafter flag pointing to the assistant model
4. Observe the crash during initialization

Workaround

Disabling speculative decoding allows the server to start:

# Remove the --drafter flag and related options
"/Applications/MLX Core.app/Contents/MacOS/mlx-serve" \
  --model ~/.omlx/models/mlx-community/gemma-4-31b-it-bf16 \
  --serve --port 8080 --host 0.0.0.0 \
  # ... (omit --drafter and draft-specific flags)

Expected Behavior

The server should initialize successfully with the drafter model loaded, and speculative decoding should accelerate token generation.

Additional Context

• The drafter model appears to load successfully (loaded 4 layers, vocab=262144, hidden=1024)
• The error occurs during the first tensor operation involving the drafter
• This may be a bug in MLX's speculative decoding implementation for Gemma-4 models with large dimension mismatches

Suggested Fixes

1. Verify Drafter Model: Check if the gemma-4-31B-it-assistant-bf16 model is correctly formatted and compatible with the main model.
2. Check MLX Version: Ensure the latest version of MLX is being used, as this may be a known issue that has been fixed.
3. Debug Tensor Operations: Add logging to identify which specific operation is producing the empty array.
4. Test with Smaller Models: Try speculative decoding with smaller Gemma models to isolate whether this is specific to the 31B scale.

Files to Investigate

• mlx-serve speculative decoding implementation
• gemma-4-31B-it-assistant-bf16 model structure
• MLX C++ backend tensor operations (ops.cpp:3701)

[ddalcu]

you are using bf16, you sure that's what you want ?
gemma-4-31b-it-bf16 -> slow, a lot of ram, a lot of disk space, no reason to use this
gemma-4-31b-it-8bit -> fast, near losses quality compared to bf16.
gemma-4-31b-it-4bit -> faster, near good quality

Drafter, yes, it's better to use the BF16 because you want 100% quality and it's very tiny so it won't saturate your memory.

I made a video using MLX-Serve to explain (some) of this 🙂
https://www.youtube.com/watch?v=kHsT8yvtn-E