Gemma 4 E2B/E4B iOS export

[Lewis300]

iOS: Gemma 4 E2B/E4B export

Adds an iOS export for the Gemma 4 per-layer-embedding variants (google/gemma-4-E2B-it,
google/gemma-4-E4B-it). Runner support is in the companion Swift PR (#302); a bundle from
this export needs both to run.

cd models/gemma4
uv run export.py --model google/gemma-4-E2B-it

produces exports/gemma_4_e2b_it_4bit_palettized_static/: the .aimodel, an INT8
per-layer-embeddings sidecar (*_ple.safetensors), metadata.json and the tokenizer.

Model (models/ios/gemma4_text.py)

• Gemma4ForCausalLMForiOS — the text decoder of the multimodal checkpoint: dual head
  dims (sliding 256 / global 512), KV shared across same-type layers, v-norm, proportional
  RoPE on the global layers, per-layer embeddings. Loads through the base from_hf
  (_HF_MODEL_CLASS = Gemma4ForConditionalGeneration); _mutate_state_dict keeps only the
  model.language_model.* keys.
• Two flat KV caches — a full-context global cache and a fixed-depth sliding-window ring
  (SLIDING_RING_SIZE, 576 for the 512 window), each holding only its layer type's storing
  layers.
• RoPE arrives precomputed as rope_cos/rope_sin rows instead of position_ids: a 131k
  position overflows a 16-bit position input, and a 32-bit one feeding the RoPE gather fails
  to compile.
• Per-layer embeddings are externalized. The multi-GB table is written as an INT8 sidecar
  and fed per step as ple_embeddings; ple_scale / ple_zero_point are frozen parameters
  the graph dequantizes with.
• The final logit soft cap is applied in the runner, not the graph, and published as
  language.final_logit_softcapping.
• Export contract hooks describe one ladder rung — build_reference_inputs /
  build_dynamic_shapes / export_static_shape_configs / export_hardware_constraints —
  with the rung's context bucket carried in TraceSpec.cache_seq_len.

Export (models/gemma4/export.py)

• Follows export/ios.py, except that each transformer (context bucket, query length)
  pair is traced by torch.export at fully static shapes and registered directly as
  extend_{ctx}_{q} / prompt_opt_{ctx}_{q}. BlockedSDPA unrolls its block loop, so the
  graph's op count depends on the context length. The gather is traced dynamic and
  specialized exactly as ios.py does it (a statically traced quantized fused dequant-gather
  doesn't convert).
• Ladder: contexts 1024 / 8192 / 32768 / 131072, up to --max-context-length (a power
  of two; default and max 131072), which becomes the top bucket. extend at q=8,
  prompt_opt at q=64 — 11 functions for the full ladder.
• Compression defaults to 4bit_palettized.yaml (4-bit per-grouped-channel k-means,
  8-bit per-tensor per-layer-embedding gate/projection); --compression none skips it.
• Fails fast: an existing bundle without --overwrite, a --max-context-length that isn't
  a power of two or is no longer than the prefill query length, or a sliding window other than
  512 fails before any weights are loaded; a checkpoint without per-layer embeddings fails while loading, before conversion.
  --overwrite only removes the old asset just before the new one is saved.
• Memory: the PLE sidecar is written straight after loading, quantized in row blocks, and
  the table is dropped before palettization.

Primitives and shared code

• BlockedSDPA (primitives/ios/sdpa.py) — online-softmax attention over the flat global
  cache in block_size chunks, with block-scaled weights so no fp16 accumulator grows with
  the key count.
• apply_rope(..., head_axis=) (primitives/ios/rope.py) — lets callers holding Q/K as
  (B, S, n_heads, head_dim) skip a transpose.
• _is_layer_key_beyond (models/base.py) — also matches keys whose
  hf_state_dict_prefix has already been stripped (layers.N. with no leading component).
• test_mistral.py — seeds TestMistraliOSAttention.test_coreai. It drew unseeded
  weights and inputs, so it passed or failed depending on which tests ran first; after the
  Gemma tests it missed its 1e-5 tolerance by ~2e-5.

Design notes

• Why a separate export.py rather than export/ios.py? Two things ios.py can't do:

  • One static graph per context bucket. BlockedSDPA unrolls its block loop, so the
    op count depends on the context length, and each rung has to be traced on its own.
    ios.py emits four fixed entrypoints into one AIProgram per call.
  • Dump the per-layer embeddings table. It is written as an INT8 sidecar next to the
    asset rather than exported into the graph, and declared in the bundle metadata.

  Otherwise export.py mirrors ios.py's _export_programs / _convert_to_coreai.

• Private helpers. export.py imports pipeline._generate_output_name and
  llm.export._load_compression_config_object to get the same bundle naming and YAML
  validation as the shared pipeline. Happy to make them public in a follow-up.

• Sliding window as a class constant. export_hardware_constraints(cls, max_context_length) receives no config, but the sliding cache's alignment depends on the
  ring depth; the export checks the checkpoint's window against it up front.

• Tied embeddings only; an untied config raises. E2B and E4B are tied.

Tests

• test_gemma4.py — chunked-prefill parity with HF over a synthetic config, including ring
  wrap and one / several / ragged global-attention blocks; the runner-side soft cap; the
  one-rung export contract; the PLE sidecar's rows and scale matching the graph's; rejection
  of unquantized PLE and non-PLE checkpoints; _is_layer_key_beyond.
• test_gemma4_export.py — a quantized tiny model through the real converter, saved and
  reloaded, with every function's name and I/O checked against what the runner looks up;
  a shortened ladder traces every rung; CLI bounds; required metadata fields.
• test_flash_attention.py — BlockedSDPA parity against SDPA and
  scaled_dot_product_attention (fp32/fp16, 65536 keys), and rank <= 4 for every op.

Testing

• pytest python/tests (with USE_LOCAL_COREAI=1): all pass.
• Export manifest (function names, input shapes, I/O names, hardware constraints) compared
  across the refactor on a synthetic model: identical.
• export.py end to end on a small local checkpoint (--compression none,
  --max-context-length 2048).
• E2B exported with the default palettization and run through llm-runner: generates.

[gokulkrishna98]

Thanks.