# Gemma 4 Assistant (MTP drafter) — tensor inventory (Stage 0)

This document records the expected tensor layout for `Gemma4AssistantForCausalLM` / `model_type: gemma4_assistant`, cross-referenced with the MLX-VLM reference implementation (PR #1112).

## HF config highlights (`google/gemma-4-26B-A4B-it-assistant`)

- `architectures`: `["Gemma4AssistantForCausalLM"]`
- `model_type`: `gemma4_assistant`
- `backbone_hidden_size`: 2816 (must match paired target Gemma 4 backbone)
- `use_ordered_embeddings`: typically `false` for 26B-A4B / 31B (dense tied LM head); `true` for E2B/E4B (centroid `MaskedEmbedder`)
- `num_centroids`: 2048 (when ordered embeddings)
- `centroid_intermediate_top_k`: 32
- Nested `text_config`: 4 layers, `layer_types` e.g. three `sliding_attention` + one `full_attention`, `attention_k_eq_v: true`, `vocab_size: 262144`, `hidden_size: 1024`, etc.

## Expected safetensors names (HF / MLX)

| Logical component | Typical HF / checkpoint name | GGUF name (this fork) |
|-------------------|------------------------------|------------------------|
| Pre-projection | `pre_projection.weight` | `mtp.pre_projection.weight` |
| Post-projection | `post_projection.weight` | `mtp.post_projection.weight` |
| Token embeddings (inner) | `model.embed_tokens.weight` | `token_embd.weight` |
| Final norm | `model.norm.weight` | `output_norm.weight` |
| Per-layer | `model.layers.{i}.*` | `blk.{i}.*` (via `tensor_mapping.py`) |
| Centroid head (E2B/E4B) | `masked_embedding.centroids.weight` | `mtp.centroids.weight` |
| Token ordering | `masked_embedding.token_ordering` | `mtp.token_ordering` (I32) |
| LM head (if untied) | `lm_head.weight` | `output.weight` |

## MLX reference files

- `mlx_vlm/speculative/drafters/gemma4_assistant/gemma4_assistant.py` — forward, `draft_block`, `sanitize`
- `mlx_vlm/speculative/drafters/gemma4_assistant/masked_embedder.py` — sparse centroid LM head
- `mlx_vlm/speculative/drafters/gemma4_assistant/masks.py` — SWA / full masks

## Notes for llama.cpp port

1. MTP consumes **target K/V** for the last sliding-attention and last full-attention layers; the assistant has **no independent KV** in the reference design. The initial C++ integration uses the assistant as a loadable arch with standard KV for bring-up; full KV sharing is wired through `--spec-type mtp` and remains dependent on target/draft pairing and cache layout.
2. `attention_k_eq_v: true` maps to missing `blk.*.attn_v.weight` (V taken from K), matching Gemma 4 handling in `gemma4-iswa.cpp`.
3. Greedy parity with the target requires byte-identical drafting; validate with `tests/test-speculative-mtp` when model paths are available.

## Centroid / ordered embeddings LM head (E2B / E4B) — implemented in MTP

When `use_ordered_embeddings=true`, the MTP step uses the masked LM head (same idea as HF `Gemma4AssistantMaskedEmbedder` / MLX `MaskedEmbedder`), implemented in `gemma4_mtp_build_one_step()` in `src/models/gemma4-assistant.cpp`:

1. `centroid_logits = mul_mat(mtp.centroids, h)` → `[n_centroids, n_tokens]`.
2. `top_k(centroid_logits, centroid_intermediate_top_k)` → centroid indices (I32).
3. `token_ordering` is viewed as `[vsc, n_centroids]` with `vsc = n_vocab / n_centroids` so each centroid column lists `vsc` token ids (matches HF flat layout); `get_rows` gathers candidate token ids for the top centroids.
4. `get_rows(token_embd, ids)` then `mul_mat` yields sparse logits over those ids.
5. Full `[n_vocab, n_tokens]` logits are built for host sampling / tracing: base tensor filled with **`-1e30f`**, then **SET_ROWS** writes the sparse logits at the selected ids (MLX-style `min(selected)-1` is optional if parity tests require it).

### GGUF / converter layout

- **`mtp.centroids.weight`**: HF checkpoint stores `[n_centroids, n_embd]` rows; numpy is written as-is so that after GGUF’s dimension packing the loader sees **`[n_embd, n_centroids]`**, matching `create_tensor(..., {n_embd, n_c})` and `mul_mat(mtp_centroids, h)` like the dense tied head.
- **`mtp.token_ordering.weight`**: stored as **I32**, length `n_vocab`.

Verification: `scripts/verify-gemma4-assistant-gguf.py` checks embedding-length parity, and when `use_ordered_embeddings` is true asserts centroid shape, I32 ordering, and `n_vocab % n_centroids == 0`.

### Tests / scripts

- Edge smoke (optional): `LLAMA_MTP_TEST_TARGET_EDGE` + `LLAMA_MTP_TEST_HEAD_EDGE` in `tests/test-speculative-mtp.cpp` (paths must exist locally; CI skips when unset).
- Server helpers: `scripts/run-gemma4-e4b-mtp-server.sh`, `scripts/run-gemma4-e2b-mtp-server.sh`.