examples/gemma4_multimodal.py: pass per_layer_inputs to decoder (#320)

## Summary

After #296 ("Move Gemma4 per-layer embeddings to embedding model") the
Gemma4 embedding sub-model emits a second output `per_layer_inputs`
(`[B, S, num_layers * per_layer_dim]`) that the decoder consumes on
every step. `examples/gemma4_multimodal.py` wasn't updated to forward
this output, so any Gemma 4 build with `hidden_size_per_layer_input > 0`
(e.g. `google/gemma-4-E2B-it`) crashes at the first `decoder.run()`:

```
ValueError: Required inputs (['per_layer_inputs']) are missing from input feed
(['inputs_embeds', ..., 'attention_mask', 'position_ids'])
```

## Fix

- `prepare_decoder_feeds()`: accept an optional `per_layer_inputs`
argument and add it to the feeds dict when present. Builds without
per-layer inputs (`hidden_size_per_layer_input == 0`, e.g. larger Gemma
4 variants) keep working — the kwarg is optional.
- `generate()`: pull `embed_out.get("per_layer_inputs")` and pass it to
`prepare_decoder_feeds` on every step.

## Verification

```
$ python examples/gemma4_multimodal.py --mode text --prompt "What is 2+2?"
...
📝  TEXT-ONLY GENERATION
================================================================
Prompt: What is 2+2?
----------------------------------------------------------------
2 + 2 = **4**
```

## Related

This is the example-side counterpart to #318 which fixed the same gap in
the L4/L5 e2e test harness.

---------

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: justinchuby

examples/gemma4_multimodal.py

@@ -227,6 +227,7 @@ def prepare_decoder_feeds(
     past_seq_len: int,
     past_kv: dict[str, np.ndarray],
     input_ids: np.ndarray,
+    per_layer_inputs: np.ndarray | None = None,
 ) -> dict[str, np.ndarray]:
     """Prepare feeds for the **decoder** session.
 
@@ -239,23 +240,34 @@ def prepare_decoder_feeds(
         past_seq_len: Number of tokens already stored in the KV cache.
         past_kv: Mapping of ``"past_key_values.{i}.key/value"`` arrays
             from the previous decoder step.
-        input_ids: ``int64[1, cur_seq_len]`` — original token ids (needed for
-            per-layer token embeddings when ``hidden_size_per_layer_input > 0``).
+        input_ids: ``int64[1, cur_seq_len]`` — original token ids (legacy
+            decoder input, retained for backward compat). On builds where
+            the embedding model emits ``per_layer_inputs`` the decoder
+            does not consume ``input_ids``.
+        per_layer_inputs: Optional ``[1, cur_seq_len, num_layers *
+            per_layer_dim]`` tensor with the same dtype as
+            ``inputs_embeds`` (that is, the model/config dtype), emitted
+            by the embedding model when
+            ``hidden_size_per_layer_input > 0``. The decoder requires
+            this input on every step (prefill + decode).
 
     Returns:
         Complete feeds dict for the decoder ONNX model.
     """
     batch_size, cur_seq_len, _ = inputs_embeds.shape
     total_seq_len = past_seq_len + cur_seq_len
 
-    return {
+    feeds: dict[str, np.ndarray] = {
         "inputs_embeds": inputs_embeds,
         # Attend to all tokens (past + current)
         "attention_mask": np.ones((batch_size, total_seq_len), dtype=np.int64),
         "position_ids": np.arange(past_seq_len, total_seq_len, dtype=np.int64)[np.newaxis, :],
         "input_ids": input_ids,
         **past_kv,
     }
+    if per_layer_inputs is not None:
+        feeds["per_layer_inputs"] = per_layer_inputs
+    return feeds
 
 
 # ---------------------------------------------------------------------------
@@ -479,10 +491,16 @@ def generate(
             )
         )
         inputs_embeds: np.ndarray = embed_out["inputs_embeds"]
+        # Gemma4 builds with hidden_size_per_layer_input > 0 also emit a
+        # second embedding output (``per_layer_inputs``) that the decoder
+        # consumes on every step. Pass it through transparently when present.
+        per_layer_inputs = embed_out.get("per_layer_inputs")
 
         # ---- Decoder session ----
         decoder_out = decoder_session.run(
-            prepare_decoder_feeds(inputs_embeds, past_seq_len, past_kv, cur_ids)
+            prepare_decoder_feeds(
+                inputs_embeds, past_seq_len, past_kv, cur_ids, per_layer_inputs
+            )
         )
 
         # Greedy: pick the token with the highest logit at the last position