Gemma4 E-series staged export should PAD-mask image-token rows for per-layer embedding inputs

## Summary

Gemma4 E-series multimodal staged exports can lose image conditioning because the transformer-stage graph uses the raw image placeholder token id when computing per-layer embedding (PLE) identity inputs.

The stage also receives already-spliced `hidden_states`, so the image features are present at the graph boundary, but the PLE side channel still encodes token identity from `image_token_id` (`258880`) instead of the PAD token id used by the Hugging Face path for multimodal/image rows. The wrong PLE identity is injected into every decoder layer and can dominate or cancel the visual conditioning.

In our fork, decoupling the image-position projection mask from the PLE token identity fixed the issue:

- image/PAD rows still use the provided `hidden_states` for the per-layer projection input,
- PLE ids for those same rows are replaced with `pad_token_id`,
- non-image text rows continue to use their original token ids.

We are filing this as an issue, not a PR, because the repository notes that PRs are closed.

## Affected Path

Observed in Gemma4 E-series staged transformer shards with `hidden_size_per_layer_input > 0`, specifically around:

- `models/macos/gemma4.py`, `Gemma4TransformerLayerStage.project_per_layer_inputs`
- the stage contract where transformer stages take both `hidden_states` and `input_ids`

The pre-fix behavior was effectively:

```python
inputs_embeds = self.embed_tokens(input_ids)
image_mask = (input_ids == image_token_id).unsqueeze(-1)
inputs_embeds = torch.where(image_mask, hidden_states, inputs_embeds)
per_layer_projection = ...
per_layer_inputs = self.get_per_layer_inputs(input_ids)
return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
```

That preserves the spliced image rows for `per_layer_projection`, but `get_per_layer_inputs(input_ids)` still sees `258880` on every image row.

The fix we validated is:

```python
inputs_embeds = self.embed_tokens(input_ids)
image_token_id = int(getattr(self.config, "image_token_id", 258880))
pad_token_id = int(getattr(self.config, "pad_token_id", 0))
image_or_pad_mask = (input_ids == image_token_id) | (input_ids == pad_token_id)
inputs_embeds = torch.where(image_or_pad_mask.unsqueeze(-1), hidden_states, inputs_embeds)

per_layer_projection = ...
per_layer_input_ids = torch.where(
image_or_pad_mask,
torch.full_like(input_ids, pad_token_id),
input_ids,
)
per_layer_inputs = self.get_per_layer_inputs(per_layer_input_ids)
return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
```

## Reproduction

Model/export shape:

- model: `google/gemma-4-E4B-it-qat-q4_0-unquantized`
- staged split: `0:21,21:42`
- native context: `131072`
- multimodal fixture: deterministic RGB image with red square, blue circle, yellow triangle, white diagonal line
- image token span: `266` rows, `image_token_id=258880`, `pad_token_id=0`
- fp32 vision tower output was validated against HF before LM execution (`0` non-finite rows, cosine `1.0`)

Symptoms before the exporter fix:

- caix/Core AI generated text like "No shapes are visible" or image-independent descriptions.
- Frozen HF soft tokens did not help, so preprocessing, vision tower, splice magnitude, and runtime insertion were ruled out.
- mm-on and mm-off teacher-forced agreement were identical in one QAT run: `43/64 == 43/64`.
- Runtime-only PAD-masking of `input_ids` improved the PLE identity but disabled the graph's image-token projection mask, so it was not sufficient.

PLE instrumentation:

- Raw image-id PLE path:
- HF-vs-CoreAI layer-0 image-row PLE cosine: `0.34706932306289673`
- all-layer raw PLE delta norm mean: about `102.109`
- PAD-masked PLE path:
- HF-vs-CoreAI layer-0 image-row PLE cosine: `1.0`
- all-layer PLE delta norm: `0.0`

After the exporter fix:

- Split-cache QAT bundle first token: `818`, matching HF.
- Teacher-forced image gate: `mm_on_hf_top1=58/64`, `mm_off_hf_top1=43/64`, first on/off divergence at step `4`.
- Live HTTP response to the fixture image: "Red square, Blue circle, Yellow triangle".
- The same runtime also preserved existing non-mm staged behavior and the Gemma 12B unified multimodal route.

## Why this looks like an exporter issue

The Hugging Face path appears to use PAD identity for multimodal/image rows in the PLE input while separately incorporating the projected image features. The current staged graph ties both decisions to the same raw `input_ids` tensor:

- image-token id is needed to decide where `hidden_states` should replace `embed_tokens(input_ids)`,
- PAD id is needed for `get_per_layer_inputs(...)` on those image rows.

Using one tensor for both roles makes the exported stage unable to match HF for E-series multimodal rows unless the graph itself decouples the masks.

The 12B `gemma4_unified` path did not expose this because its transformer stages do not use this E-series `input_ids` / PLE side channel.

## Ask

Is this a known/planned gap in the Gemma4 E-series macOS staged export path?

Would Apple consider decoupling the image-position projection mask from the PLE identity ids in `Gemma4TransformerLayerStage.project_per_layer_inputs`, so image/PAD multimodal rows use spliced `hidden_states` for projection while using `pad_token_id` for `get_per_layer_inputs`?

We can provide the minimal fixture, logs, and the validated fork patch if useful.